; analysis_sim
; Program to simulate idealised observations and background from given "truth"
; then try analysing them.
; Actual and assumed statistics are specified. 
; This version features: 
;   One type of observations, generated according to specified statistics
;   Background either a randomly perturbed or shifted version of truth.
;   Explicit calculation of matrix inverses
; RS 22 June 1998 
; Version for ASI revised 29 April 2002, corrected 24 June 2002
; ---------------------------------------------------------------------------

function truth, x, t_type, test, shift=shift
; define "truth" at an array of points (x)
; t_type=0  for truth set to 0.0
;       =1  for truth as pseudo=step function
;       =2  use sum of arbitrary sinusoidal curves
; if "shift" is set, then series is shifted (to the left) by shift units

s=size(x)
nx=s(1)
xs=x
t=fltarr(nx)

; if "shift" is specfied, add it to x
if keyword_set(shift) then begin
  if test then print,'"Truth" will be shifted by ',shift
  xs(*)=xs(*)+shift
end

if t_type eq 0 then begin
  if test then print,'Setting "truth" to zero'
  ; set t to 0
  t(*)=0.0

end else if t_type eq 1 then begin
  ; "step function" option
  ; set t to 4 at x>7, t to -4 at x<5, and smoothly varying between
  if test then print,'Setting "truth" as pseudo-step function'
  for i=0,nx-1 do begin
    if xs(i) ge 7.0 then begin
      t(i)=4.0
    end else if xs(i) le 5.0 then begin
      t(i)=-4.0 
    end else begin
      xx=0.5*(xs(i)-5.0)
      t(i)=-4.0+8.0*(3.0*xx*xx-2.0*xx*xx*xx)
    end
  end         

end else begin
  ; amplitude, frequency of sin, cos components
  if test then print,'Setting "truth" using sinusoidal variations'
  s1=2.0
  k1=1.0
  s2=0.5
  k2=2.0
  c1=0.5
  l1=0.5
  c2=0.5
  l2=2.0
  t=s1*sin(k1*xs)
  t=t+s2*sin(k2*xs)
  t=t+c1*cos(l1*xs)
  t=t+c2*cos(l2*xs)
end 

return,t
end

function frand, x, sig, L, norder, seed, test, shift=shift
; Generate a randomly varying series, at regularly spaced points x.
; The series has a nominal mean of 0.0, standard deviation of sig, and a
; correlation length L.
; The series is generated using a recursive filter technique applied to a
; random number series, which originally has a sigma of 1.0
; if shift is set, then series is shifted (to the left) by shift units

s=size(x)
nx=s(1)
dx=x(1)-x(0)
; work with extra long arrays, but only return last nx values generated
; add an extra nst points
nst=200
; if "shift" is set, amend nst
if keyword_set(shift) then begin
  nd=long(shift/dx)
  nst=nst+nd
end
nr=nx+nst
; random number array
if test then print, 'Generating random nos; nx, seed : ',nr, seed
r=fltarr(nr)
r=randomn(seed,nr)
r(0)=0.0 ; force unfiltered series to start at zero! 

; get SD of unfiltered series (nominally 1.0, with mean of 0.0)
rm=0.
rs=0.
for i=0,nr-1 do begin
  rm=rm+r(i)
  rs=rs+r(i)*r(i)
end 
rm=rm/nr
rs=rs/nr
rsd=sqrt(rs-rm*rm)
if test then print, 'Unfiltered mean, SD: ',rm, rsd

if norder ne 1 and norder ne 2 then begin
  print,'norder is not 1 or 2; it is ',norder
  stop
end  

if L gt 2.0*dx and L gt 0.0 then begin
  ; do recursive filter
  ; filter weight is derived from L and dx, assuming dx/L<<1.0
  w=1.0-dx/L
  if test then print, 'Recursive filter weight, order: ',w, norder
  for j=1,norder do begin
    for i=1,nr-1 do begin
      r(i)=w*r(i-1)+(1.0-w)*r(i)
    end
  end
  
  ; copy last nx points for output
  f=fltarr(nx)
  f=r(nr-nx:nr-1)

  ; get SD of filtered series, and rescale to requested value    
  fm=0.
  fs=0.
  for i=0,nx-1 do begin
    fm=fm+f(i)
    fs=fs+f(i)*f(i)
  end 
  fm=fm/nx
  fs=fs/nx
  fsd=sqrt(fs-fm*fm)
  if test then print, 'Filtered mean, SD: ',fm, fsd
  ; Adjust series to give requested SD (& zero mean??)
  ; try calculating scaling factor
  if norder eq 1 then begin
    wp=w
  end else begin     ; else if norder eq 2 then begin 
    wp=sqrt(w)
  end
  wfact=sig*sqrt((1.0+wp)/(1.0-wp))  ; theoretical scaling
  if test then print, 'Scaling calculated: ',wfact
  f=f-fm  ; subtract mean
end else if L gt 0.0 then begin
  print,'Correlation length is too short!',L
  stop
end else begin  ; L <= 0.0 - no filtering
  wfact=sig
  f=fltarr(nx)
  f=r(nr-nx:nr-1)
end
f=f*wfact

return,f

end
  
  
function fcst, t, x, sigf, Lf, seedf, norder, test
; generate a "forecast" background array given truth(t) and actual statistics
; both forecast and truth are valid at equally spaced grid-ponts x
; uses a simple exponential correlation model

; call frand to generate series with mean 0 and requested statistics
if sigf gt 0.0 then begin
  fc=frand(x,sigf,Lf,norder,seedf,test) 
  fc=fc+t
end else begin
  fc=t
end    

return, fc 
end

function fint, xl, f, x
; interpolate field at location xl from values f at locations x
s=size(x)
n=s(1)
if xl lt x(0) then begin 
  fi=f(0)
  return, fi
end  
if xl gt x(n-1) then begin 
  fi=f(n-1)
  return, fi
end  
for i=1,n-1 do begin
  if xl ge x(i-1) and  xl le x(i) then begin
    w=(xl-x(i-1))/(x(i)-x(i-1))
    fi=f(i-1)+w*(f(i)-f(i-1))
  end
end
return, fi
end  

function obgen, xo, sigo, Lo, seedo, base, x, norder, test
; generate a set of obs at locations xo given:
; base values (normally truth) at regularly spaced x locations
; and actual (sigma & correlation length) statistics

; get size of xo; set up ob array of same length
s=size(xo)
no=s(1)
ob=fltarr(no)
obb=fltarr(no)

if Lo gt 0.0 then begin
  ; generate obs with correlated errors
  ; first get a series at points x with specified statistics
  os=frand(x,sigo,Lo,norder,seedo,test)
  ; obs are generated from os+base
  os=os+base 
  s=size(x)
  nx=s(1)
  dx=x(1)-x(0)
  for i=0,no-1 do begin
    ; pick nearest point, rather than interpolating
    j=fix(0.5+(xo(i)-x(0))/dx)
    if j gt nx then j=nx-1
    if j lt 0 then j=0
    ob(i)=os(j)
    obb(i)=base(j)
  end
end else begin
  ; obs with UNcorrelated errors (Lo <= 0.0)
  orn=fltarr(no)
  orn=randomn(seedo,no)
  for i=0,no-1 do begin
    ; interpolate base to ob location + add error
    obb(i)=fint(xo(i),base,x)
    ob(i)=obb(i)+sigo*orn(i)
  end     
end    
if test then print,"ob x:    ",xo
if test then print,"ob base: ",obb
if test then print,"ob val:  ",ob

return, ob
end 

function ocor, x1, x2, Lo
; Define correlation between errors in two observations according to their 
; separation. 
; Use a SOAR function with length scale Lo.
; if Lo is -ve obs are assumed to be uncorrelated.
; (Warning - assumes collocated obs are UNcorrelated, if Lo=0.0, so gives
; the wrong answer for self-correlation of observation errors!)

d=abs(x1-x2)

if Lo gt 0.0 then begin
  c=(1+d/Lo)*exp(-d/Lo)
end else begin
  c=0.0
end    

return,c
end

function fcor, x1, x2, Lf
; define background error correlation between two points x1, x2
; use a SOAR function with length scale Lf

d=abs(x1-x2)

if Lf gt 0.0 then begin
  c=(1+d/Lf)*exp(-d/Lf)
end else begin
  c=0.0
  if d eq 0.0 then c=1.0
end    

return,c
end
  
; -----------------------------------------------------------------------------
; MAIN program
; -----------------------------------------------------------------------------

print,"ANALYSIS_SIM - Toy Analysis System"
print,"Constructs 1-D analysis from simulated truth, forecast and observations"
print,"with specified values and error characteristics"
print," "
print,"You may change parameters, or just press ENTER to keep existing values"
print,"To change ARRAYS you need to type in the right number of new values"


; PARAMETERS
; ++++++++++
; OPTIONS
test=0    ; 1 for extra printing
norder=2  ; order for autoregressive correlation calculations (1 or 2)
; ++++++++++
; TRUTH
t_type=0  ; 0 for truth set to 0.0
          ; 1 for truth as pseudo=step function
          ; 2 use sinusoidal curves
          ; 3 to generate "truth" statistically, rather then by specifying it
; parameters for truth generated using random numbers
; define initial covariance functions by sigma and correlation lengths
sigt=1.0
Ltr=1.0
; ++++++++++
; FORECAST
; forecast is based on "truth" shifted by "shift" units
shift=0.0
; OR perturbed according to these statistics
sigft=2.0
Lft=2.0
biasf=0.0
; statistics used by test analysis
sigf=sigft
Lf=Lft
; ++++++++++
; OBSERVATIONS
; max no. of obs
nox=11
; initial ob locations
xox=[5.0 ,6.0 ,4.0 ,7.0 ,3.0 ,8.0 ,2.0 ,9.0 ,1.0 ,10.0, 0.0]
; actual no. of observations, all one type
no=5
; Ob errors, biases; true & assumed values
sigot=1.0
Lot=0.0
biaso=0.0
sigo=sigot
Lo=Lot
; ++++++++

; set x values at which equations are solved
x0=0.
xn=10.
nx=100 ; actually have nx+1 points numbered 0 - nx
dx=(xn-x0)/nx
x=dx*findgen(nx+1)
t=fltarr(nx+1)
f=fltarr(nx+1)
a=fltarr(nx+1)
; y axis length +/-
ylen=8.0

; Start main loop
loop: print," "

; reset seeds for random no. generator
; (so random errors are reproducible with same parameters)
seedt=8.
seedo=1.
seedf=9.

; DEFINE TRUTH
print,'The truth field is one of the following:'
print,'0 - zero everywhere'
print,'1 - pseudo step function'
print,'2 - sum of several sinusoidal curves'
print,'3 - statistically generated'
print,'Current truth type: ',t_type
nnew=''
read,'New truth type: ',nnew
if keyword_set(nnew) then t_type=fix(nnew)

if t_type eq 3 then begin
  if test then print,'Generating "truth" with sigma, L: ',sigft, Lft
  t=frand(x,sigt,Ltr,norder,seedt,test) 
end else begin
  t=truth(x,t_type,test)
end  

; FORECAST STATISTICS
print,'BACKGROUND FIELD' 
print,'You may use a background field that is shifted relative to the "truth"'
print,'OR a background equal to truth + random perturbation'
print,'Current amount of background shift (+ve to left): ',shift
snew=''
read,'New shift: ',snew
if keyword_set(snew) then shift=float(snew)
print,'Current background error, sigf: ',sigf
snew=''
read,'New sigf: ',snew
if keyword_set(snew) then sigf=float(snew)
print,'Background error correlation length Lf',Lf
cnew=''
read,'New Lf (put Lf = 0.0 for uncorrelated errors): ',cnew
if keyword_set(cnew) then Lf=float(cnew)
; -----
; OPTIONALLY, add bias, and/or use incorrect statistics
; bias - zero, by default
biasf=0.0
; set actual background error the same as assumed error
sigft=sigf
Lft=Lf 
; -----

; generate background (forecast) field
if shift ne 0.0 then begin
  ; if necesary, generate second, shifted, version of "truth"
  if test then print,'Background is "truth" shifted by ',shift
  if t_type eq 3 then begin
    f=frand(x,sigt,Ltr,norder,seedt,test,shift=shift)
  end else begin
    f=truth(x,t_type,test,shift=shift)
  end
end else begin
  if test then $
  print,'Generating forecast with sigma, L, bias: ',sigft, Lft, biasf
  f=fcst(t, x, sigft, Lft, seedf, norder, test)
  f(*)=f(*)+biasf
end

; OBSERVATIONS
print,'OBSERVATIONS' 
print,'Number of obs, no: ',no
nnew=''
read,'New no (up to 11): ',nnew
if keyword_set(nnew) then no=fix(nnew)
; One set of obs with specified sigma and error correlations
fo=fltarr(no)
ob=fltarr(no)
od=fltarr(no)
xo=fltarr(no)

; locations
xo(0:no-1)=xox(0:no-1)
print,'Current ob locations',xo
change=''
read,'Do you want to change them (y or n)? ',change 
if keyword_set(change) then begin
  change=strmid(change,0,1)
  if change EQ 'y' or change EQ 'Y' then begin
    read,'new ob locations: ',xo
    xox(0:no-1)=xo(0:no-1)
  end 
end  

; observation errors
print,'Current observation error, sigo: ',sigo
snew=''
read,'New sigo: ',snew
if keyword_set(snew) then sigo=float(snew)
print,'Observation error correlation length Lo: ',Lo
cnew=''
read,'New Lo (0.0 for uncorrelated errors): ',cnew
if keyword_set(cnew) then Lo=float(cnew)
; -----
; OPTIONALLY, add bias, and/or use incorrect statistics
; bias - zero, by default
biaso=0.0
; set actual obs error the same as assumed ob error
sigot=sigo
Lot=Lo 
; -----

; Ob values
if test then print,'Generating obs with sigma, L, bias: ',$
sigot, Lot, biaso
ob(0:no-1)=obgen(xo(0:no-1),sigot,Lot,seedo,t,x,norder,test)
ob(0:no-1)=ob(0:no-1)+biaso  

; forecast fields at ob locations, O-F differences
for i=0,no-1 do begin
  fo(i)=fint(xo(i),f,x)
  od(i)=ob(i)-fo(i)
end 

; Set up M = H.Pf.Ht+R
M=fltarr(no,no)
; forecast correlations (H.Pf.Ht)
for i=0,no-1 do begin
  for j=0,no-1 do begin
    M(i,j)=sigf*sigf*fcor(xo(i),xo(j),Lf)
  end 
end 

; ob correlations R 
for i=0,no-1 do begin
  for j=0,no-1 do begin
    if i eq j then cor=1.0 else cor=ocor(xo(i),xo(j),Lo)
    M(i,j)=M(i,j)+sigo*sigo*cor
  end 
end 

; get inverse of M
Mi=invert(M,status)
if status ne 0 then begin
  print,'Status from matrix inversion =',status
  if status eq 1 then stop
end 

; evaluate analysis
for k=0,nx do begin
  xk=x(k)
  a(k)=f(k)
  for i=0,no-1 do begin
    cf=sigf*sigf*fcor(xk,xo(i),Lf)
    z=0.0
    for j=0,no-1 do begin
      z=z+Mi(i,j)*od(j)
    end 
    a(k)=a(k)+cf*z
  end
end

; plot results

title='Toy analysis (solid truth; dotted fcst; dashed analysis)'
stitle=string(sigo,Lo,sigf,Lf,$
format='("sigo=",f4.2," Lo=",f4.2,"  sigf=",f4.2," Lf=",f4.2)')
plot,x,t,yrange=[-ylen,ylen],$
title=title,subtitle=stitle 
oplot,xo(0:no-1),ob(0:no-1),psym=1
oplot,x,f,linestyle=1
oplot,x,a,linestyle=2

;stop

psfile=''
print,'If you wish to keep this plot, give the name of a Postscript file'
read,'Postscript file: ',psfile
if keyword_set(psfile) then begin
  psfile=strtrim(psfile,2)
  ; find any '.', and derive basename of file
  i = strpos(psfile,'.')
  if i ne -1 then psfile=strmid(psfile,0,i)
  print, 'postscript file: ',psfile+'.ps'        
  plopen,fn=psfile
  !p.font=0
  
  plot,x,t,yrange=[-ylen,ylen],$
  title=title, $
  subtitle=stitle
  oplot,xo(0:no-1),ob(0:no-1),psym=1
  oplot,x,f,linestyle=1
  oplot,x,a,linestyle=2

  plclose
  !p.font=-1

end

again=''
read,'Do you want to try again? (y or n): ', again
if keyword_set(again) then begin
  again=strmid(again,0,1)
  if again EQ 'y' or again EQ 'Y' then goto, loop
end 
print,'Goodbye' 

end