MODULE GradMinimizer;

IMPORT LRN, Integrator, RKF4Integrator, Minimizer; 
FROM Integrator IMPORT Time, State, Error, Velocity, Speed;
FROM Minimizer IMPORT 
  Point, Gradient, Length, Value, Slope, EvalProc, CheckProc;

IMPORT Wr, Fmt, Thread;
FROM Stdio IMPORT stderr;

REVEAL
  T = Public BRANDED OBJECT
      setSizeCalled: BOOLEAN := FALSE;
      NC: CARDINAL;           (* Number of client variables *)
      NS: CARDINAL;           (* Number of path state variables *)
      path: Path;             (* Path equations *)
      i: Integrator.T;        (* Path integrator *)
      
      (* Work areas for the "minimize" method: *)
      s0, s1: REF State;      (* Path states *)
      s0Dt: REF Velocity;     (* Path velocity *)
      se: REF Error;          (* Path integration error *)
      fxs: REF Value;         (* Function value at last "Diff" call. *)
      dfxs: REF Gradient;     (* Gradient at last evaluation *)
      wxs: REF Slope;         (* Slope (gradient norm) at last "Diff" call *)
    OVERRIDES
      name := Name;
      setPath := SetPath;
      setIntegrator := SetIntegrator;
      setSize := SetSize;
      minimize := Minimize;
      needsGradient := NeedsGradient;
    END;
    
PROCEDURE Name(m: T): TEXT =
  BEGIN
    IF m.path = NIL THEN m.path := path[1]; m.setSizeCalled := FALSE END;
    IF m.i = NIL THEN m.i := NEW(RKF4Integrator.T); m.setSizeCalled := FALSE END;
    RETURN "GradMinimizer("
     & "integrator = " & m.i.name() 
     & ", path = {" & m.path.eqs() & "}"
     & ")"
  END  Name;

PROCEDURE SetPath(m: T; p: Path): T =
  BEGIN
    m.path := p;
    m.setSizeCalled := FALSE;
    RETURN m
  END SetPath;
    
PROCEDURE SetIntegrator(m: T; i: Integrator.T): T =
  BEGIN
    m.i := i;
    m.setSizeCalled := FALSE;
    RETURN m
  END SetIntegrator;
    
PROCEDURE SetSize(m: T; n: CARDINAL): Minimizer.T =
  BEGIN
    m.setSizeCalled := TRUE;
    IF m.path = NIL THEN m.path := path[1] END;
    IF m.i = NIL THEN m.i := NEW(RKF4Integrator.T) END;

    m.NC := n;
    m.NS := m.path.size(m.NC);

    m.s0   := NEW(REF State, m.NS);
    m.s0Dt := NEW(REF State, m.NS);
    m.s1   := NEW(REF State, m.NS);
    m.se   := NEW(REF State, m.NS);
    m.fxs := NEW(REF Value);
    m.dfxs := NEW(REF Gradient, m.NC);
    m.wxs := NEW(REF Slope);
    RETURN m
  END SetSize;

PROCEDURE NeedsGradient(<*UNUSED*> m: T): BOOLEAN =
  BEGIN
    RETURN TRUE   
  END NeedsGradient;

PROCEDURE Minimize(
    m: T;
    eval: EvalProc;
    VAR x: Point;
    VAR fx: Value;
    VAR dfx: Gradient;
    dist: Length;
    tol: Length;
    check: CheckProc;
  ) =
  VAR wx: Slope; (* Slope at "x", i.e. "LRN.Norm(dfx)" *)
  VAR tx: Time;  (* Path time for "x" *)
  CONST
    ExpandFactor = 2.0d0;
    ContractFactor = 0.65155822430629449501d0; (* 2^(-1/Phi) *)
  BEGIN
    <* ASSERT m.setSizeCalled *>
    IF check # NIL THEN
      IF check(x, fx, dfx) THEN RETURN END;
    END;
    WITH 
      debug = m.debug,
      NC = m.NC,
      NS = m.NS,
      s0 = m.s0^,
      s0Dt = m.s0Dt^,
      s1 = m.s1^,
      se = m.se^,
      fxs = m.fxs^,
      dfxs = m.dfxs^,
      wxs = m.wxs^,
      dxMin = 0.1D0 * tol,
      dxMax = 10.0D0 * dist
    DO
      <* ASSERT NUMBER(x) = NC *>
      
      IF debug THEN
        Debug("dist, tol = ", LRN.T{dist, tol});
      END;

      (* Invariant: "x,fx,dfx" contain the best minimum seen so far. *)
      
      PROCEDURE PathDiff(
          READONLY dfxp: Gradient;
          wxp: Slope;
          READONLY up: LRN.T; 
          VAR xpDt: LRN.T;
          VAR upDt: LRN.T;
        ) =
        (*
          Calls "m.path.diff" and prints debugging info:
        *)
        BEGIN
          m.path.diff(dfxp, wxp, up, xpDt, upDt);
          IF debug THEN 
            IF trace THEN Trace('\n') END;
            Debug("PathDiff: u        = ", up); 
            Debug("PathDiff: velocity = ", xpDt);
            Debug("PathDiff: uDt      = ", upDt);
            PrintNewLine();
          END;
        END PathDiff;
      
      PROCEDURE CallEval(
          tq: Time;
          READONLY xq: Point; 
          VAR fxq: Value; 
          VAR dfxq: Gradient; 
          VAR wxq: Slope;
        ): BOOLEAN =
        (*
          Calls "eval"; updates "x,fx,dfx,wx" and calls "check" if
          it is a new minimum: 
        *)
        BEGIN
          IF eval(xq, fxq, dfxq) THEN RETURN TRUE END;
          wxq := LRN.Norm(dfxq);
          IF debug THEN
            IF trace THEN Trace('\n') END;
            Debug("CallEval: time     = ", LRN.T{tq}); 
            Debug("CallEval: position = ", xq); 
            Debug("CallEval: f, slope = ", LRN.T{fxq, wxq}); 
            Debug("CallEval: gradient = ", dfxq);
            PrintNewLine();
          END;
          IF fxq <= fx THEN 
            tx:= tq; x := xq; fx := fxq; dfx := dfxq; wx := wxq;
            IF check # NIL THEN
              IF check(x, fx, dfx) THEN RETURN TRUE END;
            END;
          END;
          RETURN FALSE
        END CallEval;

      PROCEDURE Diff (ts: Time; READONLY s: State; VAR sDt: Velocity) 
        RAISES {Integrator.Abort} =
        (*
          Called by "Integrate" to compute the state's derivative "sDt"
          for state "s" (see "Integrator.DiffProc".) Calls the
          client's "eval" procedure and maps it through the path
          equations.  Also saves in "wxs" the gradient magnitude and in
          "fxs" the function value.  If "fxs <= fx", updates "x" and
          "fx".
        *)
        BEGIN
          WITH
            xs = SUBARRAY(s, 0, NC),     
            us = SUBARRAY(s, NC, NS-NC),
            xsDt = SUBARRAY(sDt, 0, NC),
            usDt = SUBARRAY(sDt, NC, NS-NC)
          DO
            IF CallEval(ts, xs, fxs, dfxs, wxs) THEN RAISE Integrator.Abort END;
            PathDiff(dfxs, wxs, us, (*OUT*) xsDt, usDt)
          END
        END Diff;
      
      VAR ts0, ts1: Time;
      VAR restart: BOOLEAN;  (* TRUE if path is to be restarted at "x". *)
      VAR dxNext: Length;    (* Initial step length when restarting path *)
      VAR vs0: Speed;        (* Linear speed at time "ts0" *)
      VAR errMax: Length;    (* Maximum error allowed in step *)
      VAR trace: BOOLEAN := TRUE;

      BEGIN

        WITH
          xs0   = SUBARRAY(s0,   0,  NC),
          us0   = SUBARRAY(s0,   NC, NS-NC),
          xs0Dt = SUBARRAY(s0Dt, 0,  NC),
          us0Dt = SUBARRAY(s0Dt, NC, NS-NC),
          xs1   = SUBARRAY(s1,   0,  NC),
          xse   = SUBARRAY(se, 0, NC)
        DO
          
          tx := 0.0d0;
          wx := LRN.Norm(dfx);
          IF debug THEN
            IF trace THEN Trace('\n') END;
            Debug("IniState: time     = ", LRN.T{tx}); 
            Debug("IniState: position = ", x); 
            Debug("IniState: f, slope = ", LRN.T{fx, wx}); 
            Debug("IniState: gradient = ", dfx);
            PrintNewLine();
          END;

          restart := TRUE;
          dxNext := dist;
          errMax := 0.5d0 * tol;
          LOOP (* outer *)
            (*
              Here the path starting point is "x" at time "tx".
              The function, gradient and slope are in "fx", "dfx", wx".
            *)
            ts0 := tx; xs0 := x;
            m.path.start(dfx, wx, us0);
            IF debug THEN 
              IF trace THEN Trace('\n') END;
              Debug("reset path: s0 = ", s0);
              PrintNewLine();
            END;
            PathDiff(dfx, wx, us0, (*OUT*) xs0Dt, us0Dt);
            vs0 := LRN.Norm(xs0Dt);
            
            (* Loop until needs restarting: *)
            LOOP (* inner *)
              (* Take new step: *)
              TRY
                ts1 := ts0 + MIN(1.0d100, MAX(1.0d-100, dxNext/vs0));
                m.i.step(Diff, ts0,s0,s0Dt, ts1,s1, se);
              EXCEPT
                Integrator.Abort => RETURN
              END;

              (* Evaluate function and see if it was an improvement: *)
              IF CallEval(ts1, xs1, fxs, dfxs, wxs) THEN RETURN END;
              wxs := LRN.Norm(dfxs);
              IF tx = ts1 THEN 
                (* Good, keep going, faster: *)
                IF trace THEN Trace('\\') END;
                ts0 := ts1;
                s0 := s1;
                errMax := MIN(tol, ExpandFactor * errMax);
                dxNext := MIN(dxMax, ExpandFactor * dxNext);
              ELSIF tx # ts0 THEN
                (* Started climbing, but found a new minimum along the way. *)
                (* Restart path at current minimum: *)
                IF trace THEN Trace('U') END;
                dxNext := MAX(dxMin, ContractFactor * LRN.Dist(xs0, x));
                errMax := MAX(0.5d0 * dxMin, ContractFactor * errMax);
                EXIT (* inner *)
              ELSIF dxNext <= dxMin THEN
                (* Stricly uphill, but step was too small, so keep going on: *)
                IF trace THEN Trace('-') END;
                ts0 := ts1;
                s0 := s1;
                errMax := MIN(tol, dxMin/dxNext * errMax);
                dxNext := dxMin;
              ELSE
                (* Strictly uphill but step was large, reduce stepsize: *)
                IF trace THEN Trace('/') END;
                WITH
                  dx = LRN.Dist(xs0, xs1),
                  error = LRN.Norm(xse),
                  dxIdeal = m.i.adjustStepSize(dxNext, error, errMax, dxMin, dxMax)
                DO
                  dxNext := MAX(dxMin, MIN(dxIdeal, ContractFactor * dxNext));
                  errMax := MAX(0.5d0 * dxMin, ContractFactor * errMax);
                  IF debug THEN 
                    IF trace THEN Trace('\n') END;
                    Debug("dx     = ", LRN.T{dx});
                    Debug("error  = ", LRN.T{error});
                    Debug("dxNext = ", LRN.T{dxNext});
                  END;
                END
              END
            END (* inner *);
          END
        END
      END
    END
  END Minimize;
  
PROCEDURE Debug(msg: TEXT; READONLY x: LRN.T) =
  <* FATAL Wr.Failure, Thread.Alerted *>
  BEGIN
    Wr.PutText(stderr, "GradMinimizer ");
    Wr.PutText(stderr, msg);
    FOR i := 0 TO LAST(x) DO
      Wr.PutText(stderr, Fmt.Pad(Fmt.LongReal(x[i], Fmt.Style.Sci, 6), 15))
    END;
    Wr.PutText(stderr, "\n")
  END Debug;

PROCEDURE PrintNewLine () =
  <* FATAL Wr.Failure, Thread.Alerted *>
  BEGIN
    Wr.PutText(stderr, "\n");
  END PrintNewLine;

REVEAL 
  Path = PublicPath BRANDED OBJECT METHODS
      size(NC: CARDINAL): CARDINAL;
        (*
          Computes the number of path state variables from the number of
          function variables. 
        *)
      
      start(READONLY dfx: Gradient; wx: Slope; VAR u: LRN.T);
        (* 
          Computes the initial values of the auxiliary variables "u". *)
          
      diff(READONLY dfx: Gradient; wx: Slope; READONLY u: LRN.T; VAR xDt, uDt: LRN.T);
        (*
          Computes the state derivatives "xDt" and "uDt" from 
          the local gradient "dfx" and the auxiliary variables "u"
        *)
    END;
    

(*********************************************************************)

TYPE Path1 = Path OBJECT OVERRIDES 
      eqs := Eqs1; size := Size1; start := Start1; diff := Diff1;
    END;

PROCEDURE Eqs1(<*UNUSED*> p: Path): TEXT =
  BEGIN 
    RETURN "dx/dt = -g/|g|"
  END Eqs1;
  
PROCEDURE Size1(<*UNUSED*> p: Path; NC: CARDINAL): CARDINAL =
  BEGIN RETURN NC END Size1;
  
PROCEDURE Start1(<*UNUSED*> p: Path; 
    <*UNUSED*> READONLY dfx: Gradient; <*UNUSED*> wx: Slope; <*UNUSED*> VAR u: LRN.T
  ) =
  BEGIN END Start1;
  
PROCEDURE Diff1(<*UNUSED*> p: Path; 
    READONLY dfx: Gradient; wx: Slope;
    <*UNUSED*> READONLY u: LRN.T; VAR xDt, uDt: LRN.T
  ) =
  BEGIN
    xDt := dfx; 
    IF wx > 0.0d0 THEN LRN.Scale(-1.0d0/wx, xDt) END;
    <* ASSERT NUMBER(uDt) = 0 *>
  END Diff1;
  
(*********************************************************************)

TYPE Path2 = Path OBJECT OVERRIDES 
      eqs := Eqs2; size := Size2; start := Start2; diff := Diff2
    END;

PROCEDURE Eqs2(<*UNUSED*> p: Path): TEXT =
  BEGIN 
    RETURN "dx/dt = -g"
  END Eqs2;
  
PROCEDURE Size2(<*UNUSED*> p: Path; NC: CARDINAL): CARDINAL =
  BEGIN RETURN NC END Size2;
  
PROCEDURE Start2(
    <*UNUSED*> p: Path;
    <*UNUSED*> READONLY dfx: Gradient; 
    <*UNUSED*> wx: Slope; 
    <*UNUSED*> VAR u: LRN.T
  ) =
  BEGIN END Start2;
  
PROCEDURE Diff2(<*UNUSED*> p: Path; 
    READONLY dfx: Gradient; <*UNUSED*> wx: Slope;
    <*UNUSED*> READONLY u: LRN.T; VAR xDt, uDt: LRN.T
  ) =
  BEGIN 
    xDt := dfx; LRN.Neg(xDt);
    <* ASSERT NUMBER(uDt) = 0 *>
  END Diff2;
  
(*********************************************************************)

TYPE 
  Path3 = Path OBJECT OVERRIDES 
      eqs := Eqs3; size := Size3; start := Start3; diff := Diff3
    END;

PROCEDURE Eqs3(<*UNUSED*> p: Path): TEXT =
  BEGIN 
    RETURN 
      "dx/dt = - ((g.u)/|g|)u" & "; " &
      "du/dt = - g/|g| - ((g.u)/|g|)u"
  END Eqs3;
  
PROCEDURE Size3(<*UNUSED*> p: Path; NC: CARDINAL): CARDINAL =
  BEGIN RETURN 2*NC END Size3;
  
PROCEDURE Start3(<*UNUSED*> p: Path;
    READONLY dfx: Gradient; wx: Slope; VAR u: LRN.T
  ) =
  BEGIN 
    u := dfx;
    IF wx > 0.0d0 THEN LRN.Scale(-1.0d0/wx, u) END;
  END Start3;
  
PROCEDURE Diff3(<*UNUSED*> p: Path; 
    READONLY dfx: Gradient; wx: Slope;
    READONLY u: LRN.T; VAR xDt, uDt: LRN.T
  ) =
  BEGIN
    WITH
      a = -LRN.Dot(dfx, u)/wx, 
      b = -1.0d0/wx
    DO
      FOR i := 0 TO LAST(dfx) DO
        xDt[i] := a * u[i];
        uDt[i] := b * dfx[i] - xDt[i]
      END
    END
  END Diff3;
  
(*********************************************************************)

TYPE 
  Path4 = Path OBJECT OVERRIDES 
      eqs := Eqs4; size := Size4; start := Start4; diff := Diff4
    END;

PROCEDURE Eqs4(<*UNUSED*> p: Path): TEXT =
  BEGIN 
    RETURN 
      "dx/dt = - (1/|g|)((g.u)/(u.u))u" & "; " &
      "du/dt = - g - u"
  END Eqs4;
  
PROCEDURE Size4(<*UNUSED*> p: Path; NC: CARDINAL): CARDINAL =
  BEGIN RETURN 2*NC END Size4;
  
PROCEDURE Start4(
    <*UNUSED*> p: Path;
    READONLY dfx: Gradient;
    <*UNUSED*> wx: Slope; 
    VAR u: LRN.T
  ) =
  BEGIN u := dfx END Start4;
  
PROCEDURE Diff4(<*UNUSED*> p: Path; 
    READONLY dfx: Gradient; wx: Slope;
    READONLY u: LRN.T; VAR xDt, uDt: LRN.T
  ) =
  BEGIN
    WITH
      uu = LRN.NormSqr(u)
    DO
      IF uu = 0.0d0 THEN
        FOR i := 0 TO LAST(dfx) DO
          xDt[i] := 0.0d0;
          uDt[i] := - dfx[i]
        END
      ELSE
        WITH b = LRN.Dot(dfx, u)/(wx*uu) DO
          FOR i := 0 TO LAST(dfx) DO
            xDt[i] := - b*u[i];
            uDt[i] := - dfx[i] - u[i];
          END
        END
      END
    END
  END Diff4;

(*********************************************************************)

TYPE 
  Path5 = Path OBJECT OVERRIDES 
      eqs := Eqs5; size := Size5; start := Start5; diff := Diff5
    END;

PROCEDURE Eqs5(<*UNUSED*> p: Path): TEXT =
  BEGIN 
    RETURN 
      "dx/dt = u" & "; " &
      "du/dt = - g - |g| u"
  END Eqs5;
  
PROCEDURE Size5(<*UNUSED*> p: Path; NC: CARDINAL): CARDINAL =
  BEGIN RETURN 2*NC END Size5;
  
PROCEDURE Start5(<*UNUSED*> p: Path;
    READONLY dfx: Gradient; wx: Slope; VAR u: LRN.T
  ) =
  BEGIN 
    u := dfx; IF wx > 0.0d0 THEN LRN.Scale(-1.0d0/wx, u) END;
  END Start5;
  
PROCEDURE Diff5(<*UNUSED*> p: Path; 
    READONLY dfx: Gradient; wx: Slope;
    READONLY u: LRN.T; VAR xDt, uDt: LRN.T
  ) =
  BEGIN
    xDt := u;
    FOR i := 0 TO LAST(dfx) DO
      uDt[i] := - dfx[i] - wx * u[i]
    END
  END Diff5;
  
(*********************************************************************)

TYPE 
  Path6 = Path OBJECT OVERRIDES 
      eqs := Eqs6; size := Size6; start := Start6; diff := Diff6
    END;

PROCEDURE Eqs6(<*UNUSED*> p: Path): TEXT =
  BEGIN 
    RETURN 
      "dx/dt = -g/|g| + u/2" & "; " &
      "du/dt = -dir(g - u(g.u)/|u|^2)"
  END Eqs6;
  
PROCEDURE Size6(<*UNUSED*> p: Path; NC: CARDINAL): CARDINAL =
  BEGIN RETURN 2*NC END Size6;
  
PROCEDURE Start6(
    <*UNUSED*> p: Path;
    READONLY dfx: Gradient;
    wx: Slope;
    VAR u: LRN.T
  ) =
  BEGIN 
    u:= dfx;
    IF wx > 0.0d0 THEN LRN.Scale(-1.0d0/wx, u) END
  END Start6;
  
PROCEDURE Diff6(<*UNUSED*> p: Path; 
    READONLY dfx: Gradient; wx: Slope;
    READONLY u: LRN.T; VAR xDt, uDt: LRN.T
  ) =
  VAR dot: LONGREAL := 0.0d0;
  BEGIN
    IF wx = 0.0d0 THEN
      xDt := u;
      uDt := dfx
    ELSE
      WITH b = -1.0d0/wx DO
        FOR i := 0 TO LAST(u) DO
          uDt[i] := b*dfx[i];
          xDt[i] := uDt[i] + 0.5d0*u[i];
          dot := dot + uDt[i]*u[i];
        END;
        FOR i := 0 TO LAST(u) DO
          uDt[i] := uDt[i] - dot*u[i];
        END;
        WITH len = LRN.Norm(uDt) DO
          IF len > 0.0d0 THEN
            LRN.Scale(1.0d0/len, uDt)
          END
        END
      END
    END
  END Diff6;
  
PROCEDURE Trace(c: CHAR) =
  <* FATAL Wr.Failure, Thread.Alerted *>
  BEGIN
    Wr.PutChar(stderr, c);
    Wr.Flush(stderr)
  END Trace;

BEGIN
  path[1] := NEW(Path1);
  path[2] := NEW(Path2);
  path[3] := NEW(Path3);
  path[4] := NEW(Path4);
  path[5] := NEW(Path5);
  path[6] := NEW(Path6);
END GradMinimizer.