// This codes shows how to calculate the eigenvalue spectrum of
// the Schrodinger operator with square potential of finite depth
// and width, on a pseudo-1d domain, using SLEPc.
//
// Here, we consider the problem of a 1d quantum harmonic oscillator:
//   -0.5 \nabla^2 u + V(x)* u = E u
//  since this is an hermitian operator, its eigenvalues E are real
//  The eigenspectrum is not known in general. Yet it features few
//  qualitative properties:
//      - the lowest excitations (eigenstates with the lowest eigenvalues)
//        are bound-states (the eigenstate is localized in the well since
//        it asymptotically goes to zero as a exponential)
//        and they are quantized (discrete eigenvalues)
//      - states with higher energies (when the eigenvalue is larger than
//        the depth of the potential well) are called diffusion states.
//        they are not localized (oscillating functions). This part of
//        the spectrum is also called radiative spectrum.
//        In principle, this is a continuous spectrum.
//
//   for the default values of the potential well, there are 5 bound-states
//   before entering the diffusion spectrum.
//
//   See e.g. :
//   - Complement H-I in Quantum mechanics Vol. 1 by C. Cohen-Tannoudji, B. Diu, F. Laloe
//   - Landau & Lifshitz, Volume 3 : (Quantum Mechanics) Chap. III-22
//
// Usage :
// mpirun -np 4 FreeFem++-mpi -wg schrodinger-1d-square-well-SLEPc.edp \
//        -split 1 -npts 400 -nev 10 -sigma 0.0
//
// Authors: Julien Garaud <julien.garaud@gmail.com>
//          Pierre Jolivet <pierre.jolivet@enseeiht.fr>

/**************************************/
/*     Load PETSc & SLEPc macros      */
/**************************************/
load "PETSc"                             // PETSc plugin
macro dimension(      )2// End Of Macro  // 2D or 3D
include "macro_ddm.idp"                  // Additional DDM functions
macro def(i)i// EOM
macro init(i)i// EOM

func Pk = P1;                    // Finite-element space

/***************************************/
/*    Options for distributed solver   */
/***************************************/
int s      = getARGV("-split", 1)   ; // Refinement factor
//
int Npts   = getARGV("-npts" , 400) ; // Number of points on the perimeter
real Lx    = getARGV("-lx"   , 80.0); // Dimension of the domain
real Ly    = getARGV("-ly"   , 1)   ; // Dimension of the domain
//
int nEV    = getARGV("-nev"  , 10)  ; // Number of eigenvalues
real sigma = getARGV("-sigma", 0.0) ; // Shift
//
real width = getARGV("-width",10.0) ; // Width of the square well
real depth = getARGV("-depth", 2.0) ; // Depth of the square well
/***************************************/
/*     Verbosity and passed options    */
/***************************************/
if(verbosity > 0 && mpirank == 0) {
    cout << "********************************************"  << endl
        << " --- " << mpirank << "/" << mpisize
        << "- schrodinger-1d-square-well-SLEPc.edp "     << endl
        << "********************************************"  << endl
        << "- input parameters: "                          << endl
        << "  refinement factor  = " << s                  << endl
        << "********************************************"  << endl
        << "  nb of pts on perimeter = " << Npts
        << "  Domain size [Lx,Ly]    = ["<<Lx<<","<<Ly<<"]"<< endl
        << "********************************************"  << endl
        << "  nb of eigenvalues  = " << nEV                << endl
        << "  value of the shift = " << sigma              << endl
        << "********************************************"  << endl
        << "  width of the well = " << width               << endl
        << "  depth of the well = " << depth               << endl
        << "********************************************"  << endl
        << endl;
}
/***************************************/
/*   ###############################   */
/***************************************/
mesh Th = square(1, 1);      // Local mesh
int[int][int] intersection;   // Local-to-neighbors renumbering
real[int] D;                  // Partition of unity

/***************************************/
/*         Geometry parameters         */
/***************************************/
int[int] Labels=[1,2,3,4]; // labels : bottom, right, top, left sides

/***************************************/
/*         Finite Element space        */
/***************************************/
// Definition of the finite element space on the domain Th
// P1 are the first order Lagrange elements
fespace Vh(Th, Pk);      // local finite element space
//fespace Vh(Th,Pk,periodic=[[Labels[0],x],[Labels[2],x]]); // Periodic does not work

/***************************************/
/*            ##############           */
/***************************************/
{ // Construction of the rectangular domain
    int Xpts= int(0.5*Npts); // pts on the x-axis sides
    int Ypts= 1;             // pts on the y-axis sides

    Th = square(Xpts,Ypts,[Lx*(x-0.5),Ly*(y-0.5)],label=Labels);
    // .....
    build(Th,       // The local mesh
            s,        // Refinement factor
            intersection, // local-to-neighbors renumbering
            D,        // partition of unity
            Pk,           // FE-space
            mpiCommWorld // Communicator
         );
}
/***************************************/
/*          Problem parameters         */
/***************************************/
macro V(W,D)(0.5*(abs(x)-0.5*W < 0.0 ? 0.0 : D)) // End Of Macro

/***************************************/
/*          Problem definition         */
/***************************************/
varf vA(uh,vh)= intN(Th)        //  Definion of  the problem
( 0.5*dx(uh)*dx(vh)
  +1000*dy(uh)*dy(vh)   // cheat for pseudo-1d since periodic does not work
  //  ( 0.5*(dx(uh)*dx(vh)+dy(uh)*dy(vh))
  +( V(width,depth)-sigma)*uh*vh               )//  Bilinear form
  //  +on(Labels[3], uh = 0.0)      //  Boundary Conditions bottom
  //  +on(Labels[2], uh = 0.0)      //  Boundary Conditions right
  //  +on(Labels[1], uh = 0.0)      //  Boundary Conditions top
  //  +on(Labels[0], uh = 0.0)      //  Boundary Conditions left
;

varf vB(uh,vh)= intN(Th)  //  Definion of  the problem
( uh*vh )               //  Bilinear form
;

matrix<real> A = vA(Vh,Vh);
matrix<real> B = vB(Vh,Vh);
/***************************************/
/*      Build distributed matrices     */
/***************************************/
Mat DistA(A, intersection, D, clean = true);
Mat DistB(DistA, B, clean = true);
/***************************************/
/*          Problem resolution         */
/***************************************/
real[int]     EigenVAL(0);       // array to store eigenvalues
Vh<real>[int] def(EigenVEC)(1); // array to store eigenvectors

string ssparams =            // Parameters for the distributed EigenValue solver
  " -eps_nev " + nEV       + // Number of eigenvalues
  " -eps_type krylovschur" +
  " -eps_target "+ sigma   + // Shift value
  " -st_type sinvert "     +
  " -st_pc_type lu "       +
  " -eps_view"             +
  " -eps_gen_hermitian"      // The problem is symmetric
  ;

int k = EPSSolve
(DistA,              // matrix OP = A − sigma*B
 DistB,              //
 vectors = EigenVEC, // Array to store the FEM-EigenFunctions
 values  = EigenVAL, // Array to store the EigenValues
 sparams = ssparams  // Parameters for the distributed EigenValue solver
);

k=min(k,nEV); // some time the number of converged eigen value
              // can be greater than nev;
/***************************************/
/*           View the solution         */
/***************************************/
Vh<real> Temp;

for(int i=0;i<k;i++){
    if(!mpirank) cout << " Eigenvalue #"+i+" = "+EigenVAL[i]<<endl;
    Temp = EigenVEC[i];
    plotMPI(Th,     // The local mesh
            Temp,   // The local solution
            Pk,     // Local FE-space
            def,    // Macro for field definition
            real,   // Type
            cmm = "Psi("+i+")  EV = "+EigenVAL[i]
           )
}
