; Lorenz 63 model, with model error
;
; Stefano Migliorini (2006)
; National Centre for Earth Observation
; University of Reading
; s.migliorini@reading.ac.uk
;
; Uses a 4th-order Runge-Kutta solver

function fx,x,y,z,sigma
return,-sigma*x+sigma*y;
end

function fy,x,y,z,r
return,-x*z+r*x-y;
end

function fz,x,y,z,b
return,x*y-b*z;
end

pro lorenz, na, nb, dt, sol, moderr_flag=moderr_flag, modvar=modvar, seed=seed

if ~keyword_set(moderr_flag) then moderr_flag = 0   ; applies model error if true
if ~keyword_set(modvar) then modvar = 0.0

; checks for NaN
ind = where(~finite(sol),nind)
if nind ne 0 then message, 'provide valid initial conditions!'

; 
; for integration purposes, the increment should not be much greater than 0.01
; if dt is greater than 0.01 then dt/0.01 time steps are calculated before each output
orig_dt = dt
if dt gt 0.01 then begin
   nsteps = dt/0.01
   print, nsteps, ' intermediate steps are calculated'
   dt = 0.01
endif else nsteps = 1
   
na = long(na)
nb = long(nb)

size_sol = size(sol)
ndim = size_sol[1]
if nb ge ndim then message, 'end time nb too large wrt ndim'

; parameters
b=8.0/3.0  ; geometric factor
r=28.0     ; Rayleigh number
sigma=10.0 ; Prandl number

x = sol[na,0]
y = sol[na,1]
z = sol[na,2]

;for j=na+1,ndim-1 do begin
for j=na+1,nb do begin

  for jj=1,nsteps do begin

  ; 4th order Runge-Kutta
  f1=fx(x,y,z,sigma)
  g1=fy(x,y,z,r)
  h1=fz(x,y,z,b)
  f2=fx(x+dt*f1/2.0,y+dt*g1/2.0,z+dt*h1/2.0,sigma)
  g2=fy(x+dt*f1/2.0,y+dt*g1/2.0,z+dt*h1/2.0,r)
  h2=fz(x+dt*f1/2.0,y+dt*g1/2.0,z+dt*h1/2.0,b)
  f3=fx(x+dt*f2/2.0,y+dt*g2/2.0,z+dt*h2/2.0,sigma)
  g3=fy(x+dt*f2/2.0,y+dt*g2/2.0,z+dt*h2/2.0,r)
  h3=fz(x+dt*f2/2.0,y+dt*g2/2.0,z+dt*h2/2.0,b)
  f4=fx(x+dt*f3,y+dt*g3,z+dt*h3,sigma)
  g4=fy(x+dt*f3,y+dt*g3,z+dt*h3,r)
  h4=fz(x+dt*f3,y+dt*g3,z+dt*h3,b)
  x=x+dt*(f1+2.0*f2+2.0*f3+f4)/6.0
  y=y+dt*(g1+2.0*g2+2.0*g3+g4)/6.0
  z=z+dt*(h1+2.0*h2+2.0*h3+h4)/6.0

  endfor
  if ~moderr_flag then begin
     sol[j,0] = x
     sol[j,1] = y
     sol[j,2] = z
  endif else begin
     if n_elements(seed) eq 0 then print, 'initialising model error pseudo-random sequence'
     sol[j,0] = x + sqrt(dt)*sqrt(modvar[0])*randomn(seed)
     sol[j,1] = y + sqrt(dt)*sqrt(modvar[1])*randomn(seed)
     sol[j,2] = z + sqrt(dt)*sqrt(modvar[2])*randomn(seed)
  endelse 
endfor
dt = orig_dt

end