*********************************************************************************
*                               3-D trajectories                                *
*             Modified from Bernat Codina's 2-D trajectory script               *  
*                Jeff Duda, Iowa State University - Fall 2010                   *
*                          email: jdduda@iastate.edu                            *
*********************************************************************************

*This script will compute and plot a 3-D air parcel trajectory on a 2-D map (x by y).
*To use, you must display a field first.  Then run the script.

*Change line 24 to change the time increment to seconds to make units work.

*Set names of x-wind, y-wind, and z-wind components:
x_wind = "u"
y_wind = "v"
z_wind = "ww"

Re = 6.371e6
Pi = 3.141592653589
'q ctlinfo'
dtline = sublin(result,10)
DTword = subwrd(dtline,5)
DT = substr(DTword,1,2)
DT = DT*60
'q file'
line = sublin(result,5)
t_final = subwrd(line,12)
x_final = subwrd(line,3)
y_final = subwrd(line,6)
z_levs = subwrd(line,9)

go_on = 0
say "click a point on the map to designate a parcel"
'q pos'
x_init = subwrd(result,3)
y_init = subwrd(result,4)
'q xy2w 'x_init' 'y_init
lon_init = subwrd(result,3)
lat_init = subwrd(result,6)
if (lon_init != 'environment') 
 'q dims'
 presline = sublin(result,4)
 pres_init = subwrd(presline,6)
 timecheckline = sublin(result,5)
 t_init = subwrd(timecheckline,9)
 say "starting trajectory at lon = "lon_init" , lat = "lat_init" , pressure level = "pres_init" , and time t = "t_init
 go_on = 1 
else
 say "Must display a variable at the desired level and time to begin trajectory first. Script will now exit."
endif

if (go_on)
*Need to get list of available pressure levels
 i = 1
 while (i <= z_levs)
  'set z 'i
  'q dims'
  presline = sublin(result,4)
  pres.i = subwrd(presline,6)
  i = i + 1
 endwhile

 direction = 1
*This loop sets direction as either forward or backward and initializes the trajectory
 while(direction <= 2)
  'set lon 'lon_init
  'set lat 'lat_init
  lon1 = lon_init
  lat1 = lat_init
  lev1 = pres_init
  tt = t_init
  if (direction = 1)
   sign = 1
   say "Going forward in time"
   'set rgb 17 10 127 254'
   _col = 17
  else
   say "Going backward in time"
   sign = -1
   'set rgb 18 254 145 5'
   _col = 18
  endif
*The loop that does most of the work
  while (tt > 0 & tt < t_final+1)
   'set t 'tt
   'q dims'
   tcheckline = sublin(result,5)
   tt = subwrd(tcheckline,9)
*The if statement below obtains x-wind, y-wind, and z-wind values
   if (tt = t_init)
    'd 'z_wind
    omegaline = sublin(result,2)
    omega = subwrd(omegaline,4)
    'd 'x_wind
    uline = sublin(result,2)
    uwind = subwrd(uline,4)
    'd 'y_wind
    vline = sublin(result,2)
    vwind = subwrd(vline,4)

   else
*The below two while loops will find the pressure levels available in the control file
*that "newlev" is between.  Lev_li is the lower boundary, while lev_lf is the upper boundary.
    i = 1
    while (i <= z_levs)
     if (newlev < pres.i)
      lev_li = pres.i
     endif
     i = i + 1
    endwhile

    i = z_levs
    while (i >= 1)
     if (newlev > pres.i)
      lev_lf = pres.i
     endif
     i = i - 1
    endwhile

    delta_p = math_abs(lev_li-lev_lf)
*(log)-linearly interpolated omega, u, and v at newlev.  Uncomment the commented-out lines and comment out
*the corresponding lines to change it back to a linear interpolation
*    a = math_abs(newlev-lev_li)/delta_p
*    b = math_abs(newlev-lev_lf)/delta_p
    'define slopeomega = ('z_wind'(lev='lev_lf') - 'z_wind'(lev='lev_li')) / (log('lev_lf') - log('lev_li'))'
    'define slopeu = ('x_wind'(lev='lev_lf') - 'x_wind'(lev='lev_li')) / (log('lev_lf') - log('lev_li'))'
    'define slopev = (v(lev='lev_lf') - v(lev='lev_li')) / (log('lev_lf') - log('lev_li'))'
    'define omega = slopeomega*log('newlev') + 'z_wind'(lev='lev_lf') - (log('lev_lf')*slopeomega)'
    'define newu = slopeu*log('newlev') + 'x_wind'(lev='lev_lf') - (log('lev_lf')*slopeu)'
    'define newv = slopev*log('newlev') + 'y_wind'(lev='lev_lf') - (log('lev_lf')*slopev)'
*    'define omega = 'a'*'z_wind'(lev='lev_lf') + 'b'*'z_wind'(lev='lev_li')'
    'd omega'
    omega = subwrd(result,4)
*    'define newu = 'a'*'y_wind'(lev='lev_lf') + 'b'*'x_wind'(lev='lev_li')'
    'd newu'
    uwind = subwrd(result,4)
*    'define newv = 'a'*'y_wind'(lev='lev_lf') + 'b'*'x_wind'(lev='lev_li')'
    'd newv'
    vwind = subwrd(result,4)
   endif

   d_theta = (180/Pi)*((uwind*DT)/(Re*math_cos(lat1*(Pi/180))))
   d_phi = (180/Pi)*((vwind*DT)/Re)
*omega is in Pa/s...want it in mb/s
   dp = (omega/100)*DT
   lon2 = lon1 + sign*d_theta
   lat2 = lat1 + sign*d_phi
   lev2 = lev1 + sign*dp
*The next set of code checks to see if the parcel left the domain and stops if it does
   'q w2gr 'lon2' 'lat2
   xcheck = subwrd(result,3)
   ycheck = subwrd(result,6)
   if (xcheck > x_final | xcheck < 1 | ycheck > y_final | y_check < 1 | lev2 > pres.1 | lev2 < pres.z_levs)
    say "parcel exited domain heading for lon = "lon2", lat = "lat2", and pressure level = "lev2" at time "tt
    break
   else
    say tt' 'lon1' 'lat1' 'lev1' 'omega
   endif
*Plot mark/line for new point
   'set line '_col' 1 1'
   'q w2xy 'lon1' 'lat1
   x1 = subwrd(result,3)
   y1 = subwrd(result,6)
   'q w2xy 'lon2' 'lat2
   x2 = subwrd(result,3)
   y2 = subwrd(result,6)
   xdif = x2 - x1
   ydif = y2 - y1
   'draw mark 3 'x1' 'y1' 0.07'
   'draw mark 3 'x2' 'y2' 0.07'
   'draw line 'x1' 'y1' 'x2' 'y2
   rec1 = math_nint(lev1)
   rec2 = math_nint(lev2)
   'set string 1 c 1'
   'set strsiz 0.05 0.075'
*These if statements alter the position of the pressure text relative to the trajectory marking
*They are designed to keep the mark always to the left of the head of the line that leads to (x2,y2)
   if (xdif > 0 & ydif > 0)
    xspace = sign*-0.1
    yspace = sign*0.1
   endif
   if (xdif > 0 & ydif < 0)
    xspace = sign*0.1
    yspace = sign*0.1
   endif
   if (xdif < 0 & ydif > 0)
    xspace = sign*-0.1
    yspace = sign*-0.1
   endif
   if (xdif < 0 & ydif < 0)
    xspace = sign*0.1
    yspace = sign*-0.1
   endif
   if (do_i_draw = "yes")
*    'draw string 'x1+xspace' 'y1+yspace' 'rec1
    'draw string 'x2+xspace' 'y2+yspace' 'rec2
    do_i_draw = "no"
   else
    do_i_draw = "yes"
   endif
*Move to next point:
   'set lon 'lon2
   'set lat 'lat2
   lon1 = lon2
   lat1 = lat2
   lev1 = lev2
   newlev = lev2
   tt = tt + sign
  endwhile
  direction = direction + 1
 endwhile
endif

'set x 1 'x_final
'set y 1 'y_final
'set lev 'pres_init
'set t 't_init