
/*******************************************************************************
*
* File dfl_sap_gcr.c
*
* Copyright (C) 2007, 2011-2013 Martin Luescher
*
* This software is distributed under the terms of the GNU General Public
* License (GPL)
*
* SAP+GCR solver for the Wilson-Dirac equation with local deflation.
*
* The externally accessible functions are
*
*   double dfl_sap_gcr(int nkv,int nmx,double res,double mu,
*                      spinor_dble *eta,spinor_dble *psi,int *status)
*     Obtains an approximate solution psi of the Wilson-Dirac equation for
*     given source eta using the deflated SAP-preconditioned GCR algorithm.
*     See the notes for the explanation of the parameters of the program.
*
*   double dfl_sap_gcr2(int nkv,int nmx,double res,double mu,
*                       spinor_dble *eta,spinor_dble *psi,int *status)
*     This program calls dfl_sap_gcr() with the parameters nkv,..,status.
*     If the solver fails and status[0]=-3 or status[1]<0, the deflation
*     subspace is regenerated by calling dfl_modes(). The solver program
*     dfl_sap_gcr() is then called again and the results are passed to
*     the calling program.
*
* Depending on whether the twisted-mass flag is set or not, the programs
* solve the equation
*
*   (Dw+i*mu*gamma_5*1e)*psi=eta  or  (Dw+i*mu*gamma_5)*psi=eta
*
* respectively. The twisted-mass flag is retrieved from the parameter data
* base (see flags/lat_parms.c).

* The program dfl_sap_gcr() is based on the flexible GCR algorithm (see
* linsolv/fgcr.c). Before the solver is launched, the following parameter-
* setting programs must have been called:
*
*  set_lat_parms()        SW improvement coefficient.
*
*  set_bc_parms()         Boundary conditions and associated improvement
*                         coefficients.
*
*  set_sw_parms()         Bare quark mass.
*
*  set_sap_parms()        Parameters of the SAP preconditioner.
*
*  set_dfl_parms()        Parameters of the deflation subspace.
*
*  set_dfl_pro_parms()    Parameters used for the deflation projection.
*
* See doc/parms.pdf and the relevant files in the modules/flags directory
* for further explanations. The deflation subspace must have been properly
* initialized by the program dfl_subspace().
*
* All other parameters are passed through the argument list:
*
*  nkv       Maximal number of Krylov vectors generated before the GCR
*            algorithm is restarted.
*
*  nmx       Maximal total number of Krylov vectors that may be generated.
*
*  res       Desired maximal relative residue |eta-D*psi|/|eta| of the
*            calculated solution.
*
*  mu        Value of the twisted mass in the Dirac equation.
*
*  eta       Source field. Note that source fields must vanish at global
*            time 0 and NPR0C0*L0-1, as has to be the case for physical
*            quark fields. eta is unchanged on exit unless psi=eta (which
*            is permissible).
*
*  psi       Calculated approximate solution of the Dirac equation. psi
*            vanishes at global time 0 and NPROC0*L0-1.
*
* The argument status must point to an array of at least 2 and 3 integers
* in the case of the programs dfl_sap_gcr() and dfl_sap_gcr2(). On exit,
* the array elements contain the following values:
*
*  status[0] If the program is able to solve the Dirac equation to the
*            desired accuracy, status[0] reports the total number of Krylov
*            vectors that were required for the solution. Negative values
*            indicate that the program failed (-1: the algorithm did not
*            converge, -2: the inversion of the SW term on the odd points
*            was not safe, -3: the inversion of the diagonal parts of the
*            little Dirac operator was not safe).
*
*  status[1] Average number of GCR iterations needed for the solution of
*            the little Dirac equation in the course of the deflation
*            projection.
*
*  The program dfl_sap_gcr2() in addition returns
*
*  status[2] Average solver iteration numbers that were required for the
*            solution of the little Dirac equation when the deflation sub-
*            space had to be regenerated (if the regeneration fails, the
*            dfl_sap_gcr2() program terminates with an error message).
*
* If status[0]>=-1 and status[1]>=0, the programs return the norm of the
* residue of the calculated approximate solution. Otherwise the field psi
* is set to zero and the program returns the norm of the source eta.
*
* The SAP_BLOCKS blocks grid is automatically allocated or reallocated if
* it is not already allocated with the correct block size. The SW term is
* recalculated when needed and the gauge and SW fields are copied to the
* SAP block grid if they are not in the proper condition. Similarly, the
* little Dirac operator is updated when needed.
*
* The program dfl_sap_gcr2() can be used in place of dfl_sap_gcr() if
* some protection against the rare cases, where the little Dirac operator
* turns out to be accidentally ill-conditioned, is desired.
*
* Evidently the SAP+GCR solver is a global program that must be called on
* all processes simultaneously. The required workspaces are
*
*  spinor              2*nkv+2
*  spinor_dble         3              [2 in the case of dfl_sap_gcr()]
*  complex             2*nkv_pro+2
*  complex_dble        4
*
* (see utils/wspace.c), where nkv_pro, the maximal number of Krylov vectors
* generated before the GCR solver of the little Dirac equation is restarted,
* is a parameter set by dfl_pro_parms().
*
*******************************************************************************/

#define DFL_SAP_GCR_C

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include "mpi.h"
#include "utils.h"
#include "flags.h"
#include "block.h"
#include "uflds.h"
#include "sflds.h"
#include "linalg.h"
#include "sw_term.h"
#include "dirac.h"
#include "linsolv.h"
#include "sap.h"
#include "vflds.h"
#include "little.h"
#include "dfl.h"
#include "global.h"

static int nit,stat,nv;
static float mus;
static double mud;
static sap_parms_t spr;
static dfl_pro_parms_t dpr;


static void Dop(spinor_dble *s,spinor_dble *r)
{
   Dw_dble(mud,s,r);
}


static void Mop(int k,spinor *rho,spinor *phi,spinor *chi)
{
   int n,status;
   complex **wv;
   complex_dble **wvd;
   spinor **ws;

   wv=reserve_wv(1);
   wvd=reserve_wvd(1);
   ws=reserve_ws(1);

   dfl_s2v(rho,wv[0]);
   assign_v2vd(nv,wv[0],wvd[0]);
   ltl_gcr(dpr.nkv,dpr.nmx,dpr.res,mud,wvd[0],wvd[0],&status);
   assign_vd2v(nv,wvd[0],wv[0]);
   dfl_v2s(wv[0],ws[0]);

   Dw(mus,ws[0],chi);
   diff_s2s(VOLUME,rho,chi);
   set_s2zero(VOLUME,phi);

   for (n=0;n<spr.ncy;n++)
      sap(mus,spr.isolv,spr.nmr,phi,chi);

   diff_s2s(VOLUME,rho,chi);
   mulr_spinor_add(VOLUME,phi,ws[0],1.0f);

   if (stat>=0)
   {
      if (status>=0)
      {
         nit+=1;
         stat+=status;
      }
      else
         stat=status;
   }

   release_ws();
   release_wvd();
   release_wv();
}


double dfl_sap_gcr(int nkv,int nmx,double res,double mu,
                   spinor_dble *eta,spinor_dble *psi,int *status)
{
   int *bs,nb,isw,ifail;
   int swde,swdo,swu,swe,swo;
   double rho,rho0,fact;
   spinor **ws;
   spinor_dble **wsd,**rsd;
   dfl_parms_t dfl;

   dfl=dfl_parms();
   error_root(dfl.Ns==0,1,"dfl_sap_gcr [dfl_sap_gcr.c]",
              "Deflation parameters are not set");
   bs=dfl.bs;
   nv=dfl.Ns*VOLUME/(bs[0]*bs[1]*bs[2]*bs[3]);

   spr=sap_parms();
   error_root(spr.ncy==0,1,"dfl_sap_gcr [dfl_sap_gcr.c]",
              "SAP parameters are not set");

   dpr=dfl_pro_parms();
   error_root(dpr.nkv==0,1,"dfl_sap_gcr [dfl_sap_gcr.c]",
              "Deflation projector parameters are not set");

   blk_list(SAP_BLOCKS,&nb,&isw);

   if (nb==0)
      alloc_bgr(SAP_BLOCKS);

   if (query_grid_flags(SAP_BLOCKS,UBGR_MATCH_UD)!=1)
      assign_ud2ubgr(SAP_BLOCKS);

   if (query_flags(SWD_UP2DATE)!=1)
      sw_term(NO_PTS);

   swde=query_flags(SWD_E_INVERTED);
   swdo=query_flags(SWD_O_INVERTED);

   swu=query_grid_flags(SAP_BLOCKS,SW_UP2DATE);
   swe=query_grid_flags(SAP_BLOCKS,SW_E_INVERTED);
   swo=query_grid_flags(SAP_BLOCKS,SW_O_INVERTED);
   ifail=0;

   if (spr.isolv==0)
   {
      if ((swde==1)||(swdo==1))
         sw_term(NO_PTS);

      if ((swu!=1)||(swe==1)||(swo==1))
         assign_swd2swbgr(SAP_BLOCKS,NO_PTS);
   }
   else if (spr.isolv==1)
   {
      if ((swde!=1)&&(swdo==1))
      {
         if ((swu!=1)||(swe==1)||(swo!=1))
            assign_swd2swbgr(SAP_BLOCKS,NO_PTS);

         sw_term(NO_PTS);
      }
      else
      {
         if ((swde==1)||(swdo==1))
            sw_term(NO_PTS);

         if ((swu!=1)||(swe==1)||(swo!=1))
            ifail=assign_swd2swbgr(SAP_BLOCKS,ODD_PTS);
      }
   }
   else
      error_root(1,1,"dfl_sap_gcr [dfl_sap_gcr.c]","Unknown block solver");

   if (query_flags(U_MATCH_UD)!=1)
      assign_ud2u();

   if ((query_flags(SW_UP2DATE)!=1)||
       (query_flags(SW_E_INVERTED)==1)||(query_flags(SW_O_INVERTED)==1))
      assign_swd2sw();

   rho0=sqrt(norm_square_dble(VOLUME,1,eta));
   rho=rho0;
   status[0]=0;
   status[1]=0;

   if (ifail)
      status[0]=-2;
   else
   {
      ifail=set_Awhat(mu);

      if (ifail)
         status[0]=-3;
      else
      {
         ws=reserve_ws(2*nkv+1);
         wsd=reserve_wsd(1);
         rsd=reserve_wsd(1);

         nit=0;
         stat=0;
         mus=(float)(mu);
         mud=mu;

         fact=rho0/sqrt((double)(VOLUME)*(double)(24*NPROC));

         if (fact!=0.0)
         {
            assign_sd2sd(VOLUME,eta,rsd[0]);
            scale_dble(VOLUME,1.0/fact,rsd[0]);

            rho=fgcr(VOLUME,1,Dop,Mop,ws,wsd,nkv,nmx,res,rsd[0],psi,status);

            scale_dble(VOLUME,fact,psi);
            rho*=fact;

            if ((nit>0)&&(stat>=0))
               status[1]=(stat+nit/2)/nit;
            else if (stat<0)
               status[1]=stat;
         }
         else
         {
            rho=0.0;
            set_sd2zero(VOLUME,psi);
         }

         release_wsd();
         release_wsd();
         release_ws();
      }
   }

   if ((status[0]<-1)||(status[1]<0))
   {
      rho=rho0;
      set_sd2zero(VOLUME,psi);
   }

   return rho;
}


double dfl_sap_gcr2(int nkv,int nmx,double res,double mu,
                    spinor_dble *eta,spinor_dble *psi,int *status)
{
   double rho;
   spinor_dble **wsd;

   wsd=reserve_wsd(1);

   if (eta==psi)
   {
      assign_sd2sd(VOLUME,eta,wsd[0]);
      eta=wsd[0];
   }

   rho=dfl_sap_gcr(nkv,nmx,res,mu,eta,psi,status);

   if ((status[0]==-3)||(status[1]<0))
   {
      dfl_modes(status+2);

      error_root(status[2]<0,1,"dfl_sap_gcr2 [dfl_sap_gcr.c]",
                 "Deflation subspace regeneration failed (status = %d)",
                 status[2]);

      rho=dfl_sap_gcr(nkv,nmx,res,mu,eta,psi,status);
   }
   else
      status[2]=0;

   release_wsd();

   return rho;
}
