MODULE PZSpectrum EXPORTS Main;

(* Writes the Fourier power spectrum OF a curve. *)
(* Last edited on 2001-11-15 18:46:54 by stolfi *)

IMPORT  ParseParams, Process, FileRd, Wr, FileWr, Math;
IMPORT  OSError, Thread, Stdio, Fmt, FPut;
IMPORT PZCurve, PZLR3Chain, PZLRChain, PZFourier;

FROM Stdio IMPORT stderr;
FROM PZTypes IMPORT LONG, LONGS, NAT, BOOL;
 
<* FATAL Wr.Failure, Thread.Alerted, OSError.E *>

TYPE
  Options = RECORD
      inFile: TEXT;         (* Input curve sample file (with extension). *)
      outName: TEXT;        (* Output spectrum file name (without extension). *)
      lambdaMin: LONG;      (* Min scale length for spectrum computation. *)
      uniformize: BOOL;     (* TRUE to uniformize velocity before sampling *)
    END;
    
PROCEDURE Main() =
  BEGIN
    WITH
      o = GetOptions(),
      rp = PZLR3Chain.Read(FileRd.Open(o.inFile), headerOnly := FALSE, centralize := TRUE),
      p = rp.c^,
      m = NUMBER(p),
      c = PZCurve.New(m)
    DO 
      c.setSamples(p);
      Wr.PutText(stderr, "original period = " & FLR(c.tPeriod,8,5) & "\n");
      IF o.uniformize THEN 
        c.uniformize();
        Wr.PutText(stderr, "new period (length) = " & FLR(c.tPeriod,8,5) & "\n");
      END;
      WITH 
        pwr = ComputePowerSpectrum(c, o.lambdaMin)^
      DO
        WritePowerSpectrum(pwr, c.tPeriod, o);
      END
    END
  END Main;
  
PROCEDURE GetOptions(): Options =
  VAR o: Options;
  BEGIN
    WITH
      pp = NEW(ParseParams.T).init(stderr)
    DO
      TRY
        pp.getKeyword("-inFile");
        o.inFile := pp.getNext();
        
        pp.getKeyword("-outName");
        o.outName := pp.getNext();
        
        pp.getKeyword("-lambdaMin");
        o.lambdaMin := pp.getNextLongReal();

        o.uniformize := pp.keywordPresent("-uniformize");
      
        pp.finish();      
        
      EXCEPT                                                            
      | ParseParams.Error =>                                              
          Wr.PutText(stderr, "Usage: PZSpectrum \\\n");
          Wr.PutText(stderr, "  -inFile NAME \\\n");
          Wr.PutText(stderr, "  -outName NAME  \\\n");
          Wr.PutText(stderr, "  -lambdaMin NUM  \\\n");
          Wr.PutText(stderr, "  [ -uniformize ] \n");
        Process.Exit(1);
      END;
    END;
    RETURN o
  END GetOptions;

PROCEDURE WritePowerSpectrum(
    READONLY pwr: LONGS;
    tPeriod: LONG;
    READONLY o: Options;
  ) =
  BEGIN
    WriteFreqPower(o.outName, pwr, tPeriod);
    WriteBandPower(o.outName, pwr, tPeriod, o.lambdaMin, 8, 2.0d0)
  END WritePowerSpectrum;

PROCEDURE ComputePowerSpectrum(
    READONLY c: PZCurve.T; 
    lambdaMin: LONG;
  ): REF LONGS =
  (*
    Returns a vector "pwr[f]" with the power (square of amplitude)
    for each frequency "f", from 0 to "fMax > c.m DIV 2".
    Elements "pwr[0] and "pwr[fMax]" reflect only one term of the
    the series, all other elements include two conjugate terms. *)
  VAR p: LONG;
  BEGIN
    WITH
      m = c.m,
      estLength = 1.5d0 * c.tPeriod,
      nMin = MAX(2*m, CEILING(8.0d0*estLength/lambdaMin)),
      n = PowerOfTwo(nMin),
      fMax = n DIV 2,
      t = NEW(REF PZLRChain.T, n)^,
      u = NEW(REF PZLR3Chain.T, n)^,
      v = NEW(REF PZLR3Chain.T, n)^,
      r = NEW(REF PZLRChain.T, n)^,
      a = NEW(REF PZFourier.T, 3, n)^,
      rpwr = NEW(REF LONGS, fMax+1), pwr = rpwr^
    DO
      FOR i := 0 TO n-1 DO t[i] := c.tPeriod * FLOAT(i,LONG)/FLOAT(n,LONG) END;
      c.sample(t, u, v, r);
      PZLR3Chain.FourierTransform(u, a);
      
      WITH ax = a[0], ay = a[1], az = a[2] DO
        FOR f := 0 TO fMax DO
          WITH
            tx = ax[f], ty = ay[f], tz = ax[f]
          DO
            p := tx*tx + ty*ty + tz*tz
          END;
          IF f > 0 AND f < fMax THEN
            WITH
              tx = ax[n-f], ty = ay[n-f], tz = az[n-f]
            DO
              p := p + tx*tx + ty*ty + tz*tz
            END
          END;
          pwr[f] := p
        END
      END;
      RETURN rpwr
    END
  END ComputePowerSpectrum;

PROCEDURE PowerOfTwo(n: NAT): NAT = 
  (* Smallest power of two not less than "n" *)
  VAR m: NAT := 1;
  BEGIN
    WHILE m < n DO m := m+m END;
    RETURN m
  END PowerOfTwo;

PROCEDURE WriteFreqPower(
    name: TEXT; 
    READONLY pwr: LONGS; 
    tPeriod: LONG;
  ) =
  BEGIN  
    WITH 
      fMax = LAST(pwr),
      fileName = name & "-fpw.plt",
      wr = FileWr.Open(fileName)
    DO 
      Wr.PutText(stderr, "writing file " & fileName & " ...\n");
      Wr.PutText(wr, "# freq   lambda     rms  \n");
      Wr.PutText(wr, "# ----  ------- -------  \n");
      FOR f := 1 TO fMax DO
        WITH
          lambda = tPeriod/FLOAT(f, LONG)
        DO
          FPut.Int(wr, f, 6); 
          Wr.PutText(wr, "  ");
          FPut.LongReal(wr, lambda, 4, style := Fmt.Style.Fix); 
          Wr.PutText(wr, "  ");
          FPut.LongReal(wr, Math.sqrt(pwr[f]), 4, style := Fmt.Style.Fix); 
          Wr.PutText(wr, "\n");
        END;
      END;
      Wr.Flush(wr)
    END    
  END WriteFreqPower;

PROCEDURE WriteBandPower(
    name: TEXT; 
    READONLY pwr: LONGS;
    length: LONG;
    lambdaMin: LONG;
    nScales: NAT;
    scaleRatio: LONG;
  ) =
  VAR  
    lambda, lambdaLo, lambdaHi: LONG;
    sum: LONG;
    nb: NAT;
  BEGIN  
    WITH 
      fMax = LAST(pwr),
      fileName = name & "-bpw.plt",
      wr = FileWr.Open(fileName)
    DO 
      Wr.PutText(stderr, "writing file " & fileName & " ...\n");
      Wr.PutText(wr, "# lambdaLo lambdaHi    rms  \n");
      lambdaLo := 0.0d0;
      lambdaHi := lambdaMin;
      sum := 0.0d0;
      nb := 0;
      FOR f := fMax TO 0 BY -1 DO
        IF f = 0 THEN 
          lambda := LAST(LONG) 
        ELSE
          lambda := length/FLOAT(f, LONG)
        END;
        WHILE lambda > lambdaHi AND lambdaLo < length DO 
          FPut.LongReal(wr, lambdaLo, 4, style := Fmt.Style.Fix); 
          Wr.PutText(wr, "  ");
          FPut.LongReal(wr, lambdaHi, 4, style := Fmt.Style.Fix); 
          Wr.PutText(wr, "  ");
          FPut.LongReal(wr, Math.sqrt(sum), 4, style := Fmt.Style.Fix); 
          Wr.PutText(wr, "\n");
          INC(nb);
          lambdaLo := lambdaHi; 
          IF nb > nScales THEN 
            lambdaHi := length 
          ELSE 
            lambdaHi := scaleRatio*lambdaLo
          END;
          sum := 0.0d0;
        END;
        sum := sum + pwr[f]
      END;
      Wr.Flush(wr)
    END    
  END WriteBandPower;

PROCEDURE FLR(x: LONG; width, prec: NAT): TEXT =
  BEGIN
    RETURN Fmt.Pad(Fmt.LongReal(x, prec := prec, style := Fmt.Style.Fix), width)
  END FLR;

BEGIN
  Main()
END PZSpectrum.
(*
  Copyright © 2001 Universidade Estadual de Campinas (UNICAMP).
  Authors: Helena C. G. Leitão and Jorge Stolfi.
  
  This file can be freely distributed, used, and modified, provided
  that this copyright and authorship notice is preserved, and that any
  modified versions are clearly marked as such.
  
  This software has NO WARRANTY of correctness or applicability for
  any purpose. Neither the authors nor their employers chall be held
  responsible for any losses or damages that may result from its use.
*)