; analysis_mobs
; program to examine effects of analysing multiple idealised observations,
; given a forecast background field and assumed error covariance functions
; this version features: 
;   (up to) two different types of observations
;   explicit calculation of matrix inverses
; RS 22 June 1998
; ASI version revised 10 April 2002, amended 24 June 2002
; ---------------------------------------------------------------------------

function wf, x
; define background field at point x
; In this version wf=0.0 ALWAYS 

w=0.0

return,w
end

function ocor, x1, x2, Lo
; define correlation between two observations according to their separation
; use a SOAR function with length scale Lo
; if Lo is 0.0 obs are assumed to be uncorrelated

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 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_MOBS - Toy Analysis System"
print,"Constructs 1-D analysis from multiple idealised obs"
print,"with specified values and error characteristics"
print," "
print,"You may change the analysis parameters at the prompt"
print,"To change ARRAYS you need to type in the right number of new values"

; 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)
f=fltarr(nx+1)
a=fltarr(nx+1)
; analysis for separate ob types
at=fltarr(nx+1,2)

print," "
print,'In this programme, we use two different ob types (A and B).'
print,'The number of obs of each type should be in range 0 to 7.'
noa=3
nob=3
; max no. of each type; set initial values of ob locations
noax=7
nobx=7
; initial type A, B ob locations, values,  errors
xoa=[5.0, 4.0, 6.0, 3.0, 7.0, 2.0, 8.0]
xob=[5.0, 6.0, 4.0, 7.0, 3.0, 8.0, 2.0]
oba=1.0+fltarr(noax)
obb=-1.0+fltarr(nobx)
sigoa=2.0
Loa=0.0
sigob=1.0
Lob=2.0
; initial background  errors
sigf=1.0
Lf=1.0

; Start main loop
; SET NO. OF OBS OF EACH TYPE
loop: print,"Press ENTER to retain previous values"
print," "

; OBSERVATION VALUES & STATISTICS
; for each ob specify ob value, position 
; previous (or default) values are stored in xoa, oba, xob, obb
; BUT calculations use xo, ob arrays with type-A and type-B obs concatenated
print,'You may specify two separate observation types (A and B)'
print,'Different observation types have different error characteristics' 

; Type A obs 
print,'TYPE-A OBSERVATIONS' 
print,'Number of type A obs, noa: ',noa
nnew=''
read,'New noa (up to 7; 0 if none): ',nnew
if keyword_set(nnew) then noa=fix(nnew)
if noa gt 0 then begin
  xs=fltarr(noa)
  os=fltarr(noa)
  xs(0:noa-1)=xoa(0:noa-1)
  os(0:noa-1)=oba(0:noa-1)
  print,'Current type A ob locations',xs
  print,'Current type A ob values',os
  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 type-A locations: ',xs
      read,'new type-A values: ',os
      xoa(0:noa-1)=xs(0:noa-1)
      oba(0:noa-1)=os(0:noa-1)
    end 
    print,'New type-A ob locations',xs
    print,'New type-A ob values',os
  end 
  print,'Observation error for type A, sigoa: ',sigoa
  snew=''
  read,'New sigoa: ',snew
  if keyword_set(snew) then sigoa=float(snew)
  print,'Observation error correlation length Loa: ',Loa
  cnew=''
  read,'New Loa (0.0 for uncorrelated errors): ',cnew
  if keyword_set(cnew) then Loa=float(cnew)
  print,'sigoa, Loa: ', sigoa, Loa
end  
; type B obs uncorrelated - eg. sondes
print,'TYPE-B OBSERVATIONS' 
print,'Number of type B obs, nob: ',nob
nnew=''
read,'New nob (up to 7; 0 if none): ',nnew
if keyword_set(nnew) then nob=fix(nnew)
if nob gt 0 then begin
  xs=fltarr(nob)
  os=fltarr(nob)
  xs(0:nob-1)=xob(0:nob-1)
  os(0:nob-1)=obb(0:nob-1)
  print,'Current type-B ob locations',xs
  print,'Current type-B ob values',os
  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 type-B locations: ',xs
      read,'new type-B values: ',os
      xob(0:nob-1)=xs(0:nob-1)
      obb(0:nob-1)=os(0:nob-1)
    end 
    print,'New type-B ob locations',xs
    print,'New type-B ob values',os
  end 
  print,'Observation error for type B, sigob',sigob
  snew=''
  read,'New sigob: ',snew
  if keyword_set(snew) then sigob=float(snew)
  print,'Observation error correlation length Lob',Lob
  cnew=''
  read,'New Lob (0.0 for uncorrelated errors): ',cnew
  if keyword_set(cnew) then Lob=float(cnew)
  print,'sigob, Lob: ', sigob, Lob
end  

; no. of ob types (reset later if noa & nob are both non-zero)
nt=1
no=noa+nob
xo=fltarr(no)
fo=fltarr(no)
ob=fltarr(no)
od=fltarr(no)
; concatenate ob positions and values
if noa gt 0 then begin
  xo(0:noa-1)=xoa(0:noa-1)
  ob(0:noa-1)=oba(0:noa-1)
end   
if nob gt 0 then begin
  xo(noa:no-1)=xob(0:nob-1)
  ob(noa:no-1)=obb(0:nob-1)
end   
; analysis for individual obs (if only one ob type)
a1=fltarr(nx+1,no)

; FORECAST STATISTICS
; define covariance functions by sigma and correlation lengths
; forecast (background)
print,'BACKGROUND ("FORECAST") ERROR STATISTICS' 
print,'Background error sigf (constant)',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)

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

; Set up M = H.Pf.Ht+R, for all obs
; plus equivalents for each ob type on its own
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 (assume no correlations between obs of difft types)
; sigo & Lo are used if all obs are of 1 type
if noa gt 0 then begin
  Ma=fltarr(noa,noa)
  for i=0,noa-1 do begin
    for j=0,noa-1 do begin
      if i eq j then cor=1.0 else cor=ocor(xo(i),xo(j),Loa)
      M(i,j)=M(i,j)+sigoa*sigoa*cor
      Ma(i,j)=M(i,j)
    end 
  end 
  sigo=sigoa
  Lo=Loa
end
if nob gt 0 then begin
  Mb=fltarr(nob,nob)
  for i=0,nob-1 do begin
    for j=0,nob-1 do begin
      if i eq j then cor=1.0 else cor=ocor(xo(i+noa),xo(j+noa),Lob)
      M(i+noa,j+noa)=M(i+noa,j+noa)+sigob*sigob*cor
      Mb(i,j)=M(i+noa,j+noa)
    end 
  end 
  sigo=sigob
  Lo=Lob
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 

; calculate separate inverses for each ob type 
if noa gt 0 and nob gt 0 then begin
  nt=2
  Mai=invert(Ma,status)
  if status ne 0 then begin
    print,'Status from matrix inversion (ob type a) =',status
    if status eq 1 then stop
  end 
  Mbi=invert(Mb,status)
  if status ne 0 then begin
    print,'Status from matrix inversion (ob type b) =',status
    if status eq 1 then stop
  end 
end
  
; For "single ob" analysis wt for 1 ob
W1=1.0/(sigf*sigf+sigo*sigo) 

; evaluate analysis
for k=0,nx do begin
  xk=x(k)
  f(k)=wf(xk)
  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
    if nt eq 1 then a1(k,i)=f(k)+cf*W1*od(i)
  end
  if nt eq 2 then begin
    at(k,0)=f(k)
    for i=0,noa-1 do begin
      cf=sigf*sigf*fcor(xk,xo(i),Lf)
      z=0.0
      for j=0,noa-1 do begin
        z=z+Mai(i,j)*od(j)
      end
      at(k,0)=at(k,0)+cf*z
    end
    at(k,1)=f(k)
    for i=0,nob-1 do begin
      cf=sigf*sigf*fcor(xk,xo(noa+i),Lf)
      z=0.0
      for j=0,nob-1 do begin
        z=z+Mbi(i,j)*od(j+noa)
      end
      at(k,1)=at(k,1)+cf*z
    end
  end
end

; plot results

if nt eq 1 then begin 
  title='Toy analysis (dotted individual obs; dashed all obs together)'
  stitle=string(sigo,Lo,sigf,Lf,$
  format='("sigo=",f4.2," Lo=",f4.2,'+$
  '"  sigf=",f4.2," Lf=",f4.2)')
end else begin
  title='Toy analysis (dotted 1 ob type; dashed all obs)'
  stitle=string(sigoa,Loa,sigob,Lob,sigf,Lf,$
  format='("+ sigoa=",f4.2," Loa=",f4.2,"  x sigob=",f4.2," Lob=",f4.2,'+$
  '"  sigf=",f4.2," Lf=",f4.2)')
end
plot,x,f,yrange=[-2.,2.],$
title=title, $
subtitle=stitle
if noa gt 0 then oplot,xo(0:noa-1),ob(0:noa-1),psym=1
if nob gt 0 then oplot,xo(noa:noa+nob-1),ob(noa:noa+nob-1),psym=7
if nt eq 1 then begin
  for i=0,no-1 do begin
    oplot,x,a1(*,i),linestyle=1
  end
end else for i=0,nt-1 do begin
  oplot,x,at(*,i),linestyle=1
end    
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,f,yrange=[-2.,2.],$
  title=title, $
  subtitle=stitle
  if noa gt 0 then oplot,xo(0:noa-1),ob(0:noa-1),psym=1
  if nob gt 0 then oplot,xo(noa:noa+nob-1),ob(noa:noa+nob-1),psym=7
  if nt eq 1 then begin
    for i=0,no-1 do begin
      oplot,x,a1(*,i),linestyle=1
    end
  end else for i=0,nt-1 do begin
    oplot,x,at(*,i),linestyle=1
  end    
  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 begin
    print," "
    goto, loop
  end   
end 
print,'Goodbye' 

end