MODULE SPUniSolve EXPORTS Main;

(* Solves the Helmholtz differential equation by uniscale iteration. *)

(*
  This program solves the Helmholz differential equation on the sphere
  
    SLapl(u)(x) + c*u(x) = F(x, u(x))       (1)
    
  where "SLapl" is the spherical Laplacian operator, "c" is a given
  constant, "F" is a given function, and "u" is the function to be
  determined.
  
  It is assumed that an approximate solution "u" to (1) in the
  corresponding function space "V" can be found by solving the
  non-linear system
  
    H a = b         (2)
    
  where 
  
    "a[i]" is coefficient "i" of the solution 
           "u" (to be determined) in the basis "bas"; 
  
    "b[i] = <F | bas[i]>, where "F" is the spherical function
            "x -> F(x,u(x))" and "<|>" is the scalar product.
           
    "H[i,j] = <SGrad(bas[i]) | SGrad(bas[j]> + c*<bas[i]|bas[j]>
            where "SGrad" is the spherical gradient operator.
        
  Note that this system is not trivial since the right-hand side "b"
  depends on the function "u", and therefore on the coefficient vector
  "a", usually in a non-linear way.  So (2) must be solved
  iteratively.
  
*)

IMPORT Wr, Math, Thread, ParseParams, Fmt, Process;
IMPORT HLR3, LR4;
FROM SPFile IMPORT OpenRead, OpenWrite;
IMPORT SPTriang, SPMatrix, SPPWFunction, SPNLFunction; 
FROM SPLinearSystem IMPORT 
  GuessSol, 
  CholeskiSolve, ConjugateGradientSolve, GaussSeidelIteration,
  BuildFunction, UpdateSolution, PlotFunctionAndError,
  PrintMaxErrorValues;
IMPORT SPAnalyticPWFunction;
FROM SPPWFunction IMPORT Id, Basis;
FROM Math IMPORT sqrt; 

FROM Stdio IMPORT stderr;

TYPE 
  Options = RECORD
      triName: TEXT;       (* Name of triangulation file (minus extension). *)
      rhsName: TEXT;       (* Name of right-hand-side function. *)
      solName: TEXT;       (* Name of correct solution function. *)
      baseName: TEXT;      (* Name of basis file (minus extension). *)
      outName: TEXT;       (* Prefix for output file names. *)
      matName: TEXT;       (* Prefix for names of  matrix files. *)
      maxIter: CARDINAL;   (* Maximum iterations *)
      relTol: LONGREAL;    (* Max change in "u" coeffs between iterations. *)
      absTol: LONGREAL;    (* Max error compared to true sol. *)
      order: CARDINAL;     (* *)
      (* Method used to solve the linear system: *)
      choleski: BOOL;      (* TRUE uses direct method based on Choleski factorization. *)
      gaussSeidel: BOOL;   (* TRUE uses a single Gauss-Seidel iteration. *)
      conjGrad: BOOL;      (* TRUE uses the conjugate gradient method. *)
      omega: LONG;         (* Overrrelaxation factor for Gauss-Seidel. *)
      (* Plotting options: *)
      plot: BOOL;          (* TRUE to generate function and error plots. *)
      obs: HLR3.Point;     (* Position of observer *)
      gridSize: CARDINAL;  (* Number of grid steps in Lat & Lon *)
    END;

TYPE
  BOOL = BOOLEAN;
  LONG = LONGREAL;
  NAT = CARDINAL;
  
<* FATAL Wr.Failure, Thread.Alerted *>

PROCEDURE Main() =
  VAR HL, H: SPMatrix.T;
      d: LONG;
      iter: CARDINAL := 0;
      VAR uf: SPPWFunction.T;
      fMax, eMax: LONG;
  BEGIN
    WITH 
      o = GetOptions(),
      rdtri = OpenRead(o.triName & ".tri"),
      tri = MakeTriRef(SPTriang.Read(rdtri)),
      
      rdbas = OpenRead(o.baseName & ".bas"),
      basis = SPPWFunction.ReadBasis(rdbas, tri)^,
      dim = NUMBER(basis),
      
      F = MakeRightHandSide(o.rhsName),   (* Right-hand side *)
      
      ufCorr = MakeCorrectSolution(o.solName, tri), (* Correct solution *)
      
      u = NEW(REF ARRAY OF LONG, dim)^,
      b = NEW(REF ARRAY OF LONG, dim)^,
      y = NEW(REF ARRAY OF LONG, dim)^,
      v = NEW(REF ARRAY OF LONG, dim)^,
      errWr = OpenWrite(o.outName & ".erp")
    DO  
      IF o.choleski THEN 
        WITH 
          rd = OpenRead(o.matName & "-HL.mat")
        DO
          SPMatrix.Read(rd, HL);
        END
      ELSE
        WITH 
          rd = OpenRead(o.matName & "-H.mat")
        DO
          SPMatrix.Read(rd, H);
        END
      END;
      
      BEGIN 
        WITH solName = o.outName & "-0000" DO
          GuessSol(u);
          BuildFunction(basis, u, solName, ufCorr, uf);
        END;
        LOOP
          PrintMaxErrorValues(uf, ufCorr, fMax, eMax);
          Wr.PutText(errWr, FI(iter, 6) & " " & FLR(eMax, 12, 8) & "\n");
          IF (iter > 0 AND d <= o.relTol)
          OR (eMax <= o.absTol)
          OR iter >= o.maxIter THEN EXIT END;
          WITH
            solName = o.outName & "-" & Fmt.Pad(Fmt.Int(iter), 4, '0')
          DO
            ComputeRightHandSide(tri, basis, uf, F, o.order, b);
            IF o.choleski THEN
              CholeskiSolve(HL, b, y, v)
            ELSIF o.gaussSeidel THEN
              GaussSeidelIteration(H, b, o.omega, u, v)
            ELSIF o.conjGrad THEN
              ConjugateGradientSolve(H, b, u, v)
            ELSE
              <* ASSERT FALSE *>
            END;
            d := UpdateSolution(v, u);
            BuildFunction(basis, u, solName, ufCorr, uf);
          END;
          INC(iter);
        END;
        Wr.Close(errWr);
        IF o.plot THEN
          PlotFunctionAndError(o.outName, uf, ufCorr, tri, fMax, eMax, o.obs, o.gridSize);
        END;
      END;  
    END;  
  END Main;
  
PROCEDURE ComputeRightHandSide(
    tri: REF SPTriang.T;      (* Triangulation to use for integration. *)
    READONLY bas: Basis;      (* Basis for the approximation space. *)
    uf: SPPWFunction.T;       (* Current solution. *)
    F: SPNLFunction.T;        (* Right-hand side of differential eqn. *)
    sOrder: NAT;              (* Order of sampling grid in each triangle *)
    (* Input/output: *)
    VAR b: ARRAY OF LONG;     (* Right-hand-side of (2). *) 
  ) =
  (*
    Computes the right-hand side of the system (2) for the given
    solution  "uf" which has coordinates "u" relative to basis "bas".
  *)
  VAR sb2: LONG;
  BEGIN
    WITH
      dim = NUMBER(bas)
    DO
      Wr.PutText(stderr, "computing right-hand side \"b\"\n");
      sb2 := 0.0d0;
      FOR i := 0 TO dim-1 DO
        Wr.PutText(stderr, ".");
        WITH 
          bNew = SPPWFunction.DotProduct(uf, bas[i], sOrder, F, Id, tri),
          db = b[i] - bNew
        DO
          sb2 := sb2 + db*db;
          b[i] := bNew;
        END;
        (* Wr.PutText(stderr, "  b[" & Fmt.Pad(Fmt.Int(i),3) & "] = "); 
         Wr.PutText(stderr, 
          Fmt.Pad(Fmt.LongReal(b[i], style := Fmt.Style.Fix, prec := 6), 9)
          & "\n"
        ); *)
      END;
      Wr.PutText(stderr, "\n");
      Wr.PutText(stderr, " change in b =" & 
          Fmt.Pad(Fmt.LongReal(sqrt(sb2), style := Fmt.Style.Fix, prec := 6), 9)
          & "\n"
      );
    END;
  END ComputeRightHandSide;

PROCEDURE MakeTriRef(READONLY tri: SPTriang.T): REF SPTriang.T = 
  (*
    Droga de Modula-3... *)
  BEGIN
    WITH t = NEW(REF SPTriang.T) DO t^ := tri; RETURN t END
  END MakeTriRef;

PROCEDURE MakeCorrectSolution(
    name: TEXT; 
    tri: REF SPTriang.T;
  ): SPPWFunction.T =
  BEGIN
    WITH f = SPAnalyticPWFunction.FromName(name, tri) DO
      IF f = NIL THEN 
        Wr.PutText(stderr, "unrecognized function name\"" & name & "\"\n");
        Process.Exit(1)
      END;
      RETURN f
    END;
  END MakeCorrectSolution;
  
PROCEDURE MakeRightHandSide(name: TEXT): SPNLFunction.T =
  BEGIN
    WITH f = SPNLFunction.FromName(name) DO
      IF f = NIL THEN 
        Wr.PutText(stderr, "unrecognized RHS name\"" & name & "\"\n");
        Process.Exit(1)
      END;
      RETURN f
    END
  END MakeRightHandSide;

PROCEDURE GetOptions (): Options =
  VAR o: Options;
  BEGIN
    TRY
      WITH  pp = NEW(ParseParams.T).init(stderr) DO 
      
        pp.getKeyword("-triName");  
        o.triName := pp.getNext();
        
        pp.getKeyword("-baseName");                               
        o.baseName := pp.getNext();  
       
        pp.getKeyword("-rhsName");                               
        o.rhsName := pp.getNext();  
       
        pp.getKeyword("-solName");                               
        o.solName := pp.getNext();  
       
        pp.getKeyword("-outName");                               
        o.outName := pp.getNext();  
                                           
        o.choleski :=    pp.keywordPresent("-choleski");
        o.gaussSeidel := pp.keywordPresent("-gaussSeidel");
        IF o.gaussSeidel AND pp.keywordPresent("-omega") THEN
          o.omega := pp.getNextLongReal(0.0d0, LAST(LONG))
        END;
        o.conjGrad :=    pp.keywordPresent("-conjGrad");
        IF NOT(o.choleski OR o.gaussSeidel OR o.conjGrad) THEN
          pp.error("must specify a linear solution method")
        ELSIF (o.choleski AND o.gaussSeidel) 
        OR    (o.choleski AND o.conjGrad)
        OR    (o.gaussSeidel AND o.conjGrad) THEN
          pp.error("please specify only one linear solution method")
        END;
        
        pp.getKeyword("-matName");                               
        o.matName := pp.getNext();  
                  
        IF pp.keywordPresent("-relTol") THEN
          o.relTol := pp.getNextLongReal(0.0d0,LAST(LONGREAL))
        ELSE
          o.relTol := 0.0d0;
        END;
      
        IF pp.keywordPresent("-absTol") THEN
          o.absTol := pp.getNextLongReal(0.0d0,LAST(LONGREAL))
        ELSE
          o.absTol := 0.0d0;
        END;
      
        IF pp.keywordPresent("-maxIter") THEN
          o.maxIter := pp.getNextInt(0, LAST(CARDINAL))
        ELSE
          o.maxIter := 20   (* LAST(CARDINAL) *)
        END;
      
        pp.getKeyword("-order");
        o.order := pp.getNextInt(0,LAST(CARDINAL));
        
        (* Plotting options: *)
        
        o.plot := pp.keywordPresent("-plot");
        
        IF pp.keywordPresent("-obs") THEN
          o.obs := ParsePointOption(pp)
        ELSE
          (* Default is all zeros, meaning rotate sphere for best view. *)
          o.obs := HLR3.Point{c := LR4.T{0.0d0, 0.0d0, 0.0d0, 0.0d0}}
        END;

        IF pp.keywordPresent("-gridSize") THEN
          o.gridSize := pp.getNextInt(5, 1000)
        ELSE
          o.gridSize := 100
        END;

        pp.finish();
      END
    EXCEPT                                                            
    | ParseParams.Error =>                                              
        Wr.PutText(stderr, "Usage: SPUniSolve \\\n");
        Wr.PutText(stderr, "  -triName <triName>  \\\n");
        Wr.PutText(stderr, "  -baseName <baseName> \\\n");
        Wr.PutText(stderr, "  -rhsName <functionName> \\\n");
        Wr.PutText(stderr, "  -solName <functionName> \\\n");
        Wr.PutText(stderr, "  -outName <name> \\\n");
        Wr.PutText(stderr, "  [ -choleski | \\\n");
        Wr.PutText(stderr, "    -gaussSeidel [-omega NUM] | \\\n");
        Wr.PutText(stderr, "    -conjGrad ] \\\n");
        Wr.PutText(stderr, "  -matName <matName> \\\n");
        Wr.PutText(stderr, "  [ -relTol <num> ] \\\n");
        Wr.PutText(stderr, "  [ -absTol <num> ] \\\n");
        Wr.PutText(stderr, "  [ -maxIter <num> ] \\\n");
        Wr.PutText(stderr, "  -order <num> \\\n");
        Wr.PutText(stderr, "  [ -plot ]\\\n");
        Wr.PutText(stderr, "  [ -obs OW OX OY OZ ] \\\n");
        Wr.PutText(stderr, "  [ -gridSize NUM ]  \n");
        Process.Exit (1);
    END;
    RETURN o
  END GetOptions;
  
PROCEDURE ParsePointOption(pp: ParseParams.T): HLR3.Point RAISES {ParseParams.Error} =
  VAR p: HLR3.Point;
  BEGIN
    FOR i := 0 TO 3 DO
      p.c[i] := pp.getNextLongReal(-1.0d20, +1.0d20)
    END;
    RETURN p
  END ParsePointOption;  

PROCEDURE FLR(x: LONGREAL; wid, prec: CARDINAL): TEXT =
  BEGIN
    RETURN Fmt.Pad(Fmt.LongReal(x, style:= Fmt.Style.Sci, prec:= prec), wid)
  END FLR;
  
PROCEDURE FI(x: INTEGER; wid: CARDINAL): TEXT =
  BEGIN
    RETURN Fmt.Pad(Fmt.Int(x), wid)
  END FI;

BEGIN
  Main();
END SPUniSolve.