NCL geostrophic

;This is a NCL file with custom functions for the GFI dynamic meteorology group


load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
load "WRFUserARW_new.ncl"



;a function for calculating the geostrophic and ageostrophic wind along and across a given cross section

undef("wrf_geostrophic")
function wrf_geostrophic(nc_file:file,it:numeric,lat_diff_f:numeric,lon_diff_f:numeric, \
                         WE:logical,NS:logical)

begin
a=nc_file
  P_total        = wrf_user_getvar(a,"pres",it)      
  P_base         = wrf_user_getvar(a,"PB",it)      
  G_pert         = wrf_user_getvar(a,"PH",it)
  G_base         = wrf_user_getvar(a,"PHB",it)
  q_vapor        = wrf_user_getvar(a,"QVAPOR",it)
  q_cloud        = wrf_user_getvar(a,"QCLOUD",it)
  q_rain         = wrf_user_getvar(a,"QRAIN",it)
  mu             = wrf_user_getvar(a,"MU",it)
  mub            = wrf_user_getvar(a,"MUB",it)
  eta_full       = wrf_user_getvar(a,"ZNW",it)
  eta_half       = wrf_user_getvar(a,"ZNU",it)
  f              = wrf_user_getvar(a,"F",it)
  ;P_total        = P_pert+P_base
  G_total        = G_pert+G_base
  mu_total       = mu+mub
  delta_y        = 3000.
  delta_x        = 3000.
  dimG           = dimsizes(G_pert)
  dimP           = dimsizes(P_total)
;;;;;;Calculating all finite differences;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  delta_G_eta    = G_total(1:dimG(0)-1,:,:)-G_total(0:dimG(0)-2,:,:)
  delta_G_x      = G_total(:,:,1:dimG(2)-1)-G_total(:,:,0:dimG(2)-2)
  delta_G_y      = G_total(:,1:dimG(1)-1,:)-G_total(:,0:dimG(1)-2,:)
  delta_P_eta    = P_total(1:dimP(0)-1,:,:)-P_total(0:dimP(0)-2,:,:)
  delta_P_x      = P_total(:,:,1:dimP(2)-1)-P_total(:,:,0:dimP(2)-2)
  delta_P_y      = P_total(:,1:dimP(1)-1,:)-P_total(:,0:dimP(1)-2,:)
  delta_eta_full = eta_full(1:dimG(0)-1)-eta_full(0:dimG(0)-2)
  delta_eta_half = eta_half(1:dimP(0)-1)-eta_half(0:dimP(0)-2)
  
  delta_eta_full_conform_P     = conform_dims((/dimP(0),dimP(1),dimP(2)/),delta_eta_full,0)
  delta_eta_half_conform_G     = conform_dims((/dimG(0)-2,dimG(1),dimG(2)/),delta_eta_half,0)
  delta_G_eta                  = delta_G_eta/delta_eta_full_conform_P
  delta_P_eta                  = delta_P_eta/delta_eta_half_conform_G
;;;;;;;;Calculating alpha and alpha_D;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  mu_total_conform_P        = conform_dims((/dimP(0),dimP(1),dimP(2)/),mu_total,(/1,2/))
  alpha_D                   = -delta_G_eta/mu_total_conform_P 
  alpha                     = alpha_D*(1+q_vapor+q_cloud+q_rain)
  
;;;;;;;;getting right dimensions and interpolating to vg;;;;;;;;;;;;;  

  mu_total_interpol_x           = (mu_total(:,0:dimP(2)-2)+mu_total(:,1:dimP(2)-1))/2 
  mu_total_interpol_y           = (mu_total(0:dimP(1)-2,:)+mu_total(1:dimP(1)-1,:))/2
  alpha_interpol_x              = (alpha(:,:,0:dimP(2)-2)+alpha(:,:,1:dimP(2)-1))/2
  alpha_interpol_y              = (alpha(:,0:dimP(1)-2,:)+alpha(:,1:dimP(1)-1,:))/2
  alpha_D_interpol_x            = (alpha_D(:,:,0:dimP(2)-2)+alpha_D(:,:,1:dimP(2)-1))/2
  alpha_D_interpol_y            = (alpha_D(:,0:dimP(1)-2,:)+alpha_D(:,1:dimP(1)-1,:))/2
  f_interpol_x                  = (f(:,0:dimP(2)-2)+f(:,1:dimP(2)-1))/2
  f_interpol_y                  = (f(0:dimP(1)-2,:)+f(1:dimP(1)-1,:))/2
  
  delta_P_eta_interpol_vg       = (delta_P_eta(0:dimG(0)-4,:,0:dimG(2)-2)+delta_P_eta(1:dimG(0)-3,:,0:dimG(2)-2)+delta_P_eta(1:dimG(0)-3,:,1:dimG(2)-1)+delta_P_eta(0:dimG(0)-4,:,1:dimG(2)-1))/4
  delta_P_eta_interpol_ug       = (delta_P_eta(0:dimG(0)-4,0:dimG(1)-2,:)+delta_P_eta(1:dimG(0)-3,0:dimG(1)-2,:)+delta_P_eta(1:dimG(0)-3,1:dimG(1)-1,:)+delta_P_eta(0:dimG(0)-4,1:dimG(1)-1,:))/4
  
  delta_G_x_interpol_vg         = (delta_G_x(1:dimG(0)-3,:,0:dimG(2)-2)+delta_G_x(2:dimG(0)-2,:,0:dimG(2)-2))/2
  delta_G_y_interpol_ug         = (delta_G_y(1:dimG(0)-3,0:dimG(1)-2,:)+delta_G_y(2:dimG(0)-2,0:dimG(1)-2,:))/2

  mu_total_interpol_x_conform_P = conform_dims((/dimP(0),dimP(1),dimP(2)-1/),mu_total_interpol_x,(/1,2/))
  mu_total_interpol_y_conform_P = conform_dims((/dimP(0),dimP(1)-1,dimP(2)/),mu_total_interpol_y,(/1,2/))
  f_interpol_x_conform_P        = conform_dims((/dimP(0),dimP(1),dimP(2)-1/),f_interpol_x,(/1,2/))
  f_interpol_y_conform_P        = conform_dims((/dimP(0),dimP(1)-1,dimP(2)/),f_interpol_y,(/1,2/))
  delta_eta_half_conform_Gx     = conform_dims((/dimG(0)-2,dimG(1),dimG(2)-1/),delta_eta_half,0)
  delta_eta_half_conform_Gy     = conform_dims((/dimG(0)-2,dimG(1)-1,dimG(2)/),delta_eta_half,0)

;;;;;;split it up in two terms and adding afterwards;;;;;;;;;;;;;;;;;
vg_term1 = mu_total_interpol_x_conform_P(1:dimP(0)-2,:,:)*alpha_interpol_x(1:dimG(0)-3,:,:)/f_interpol_x_conform_P(1:dimP(0)-2,:,:)*delta_P_x(1:dimP(0)-2,:,:)/delta_x

ug_term1 = -mu_total_interpol_y_conform_P(1:dimP(0)-2,:,:)*alpha_interpol_y(1:dimG(0)-3,:,:)/f_interpol_y_conform_P(1:dimP(0)-2,:,:)*delta_P_y(1:dimP(0)-2,:,:)/delta_y

vg_term2 = delta_P_eta_interpol_vg*delta_G_x_interpol_vg/delta_x*alpha_interpol_x(1:dimG(0)-3,:,:)/(alpha_D_interpol_x(1:dimG(0)-3,:,:)*f_interpol_x_conform_P(1:dimP(0)-2,:,:))


ug_term2 = -delta_P_eta_interpol_ug*delta_G_y_interpol_ug/delta_y*alpha_interpol_y(1:dimG(0)-3,:,:)/(alpha_D_interpol_y(1:dimG(0)-3,:,:)*f_interpol_y_conform_P(1:dimP(0)-2,:,:))

vg=(vg_term1+vg_term2)/mu_total_interpol_x_conform_P(1:dimP(0)-2,:,:)
ug=(ug_term1+ug_term2)/mu_total_interpol_y_conform_P(1:dimP(0)-2,:,:)

;interpolate the geostrophic winds to mass points
vga=(vg(:,:,0:dimP(2)-3)+vg(:,:,1:dimP(2)-2))/2
uga=(ug(:,0:dimP(1)-3,:)+ug(:,1:dimP(1)-2,:))/2
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;Calculate the angle between th cross section and the regular grid
 if(WE .and. lat_diff_f.ne.0.) then
    gamma            = atan(lat_diff_f/lon_diff_f)
 else if(WE .and. lat_diff_f.eq.0.) then
        gamma=0.
      end if
 end if  

 if(NS .and. lon_diff_f.lt.0.)then
   lat_diff_f=-lat_diff_f
   gamma=3.14/2.0-atan(lat_diff_f/lon_diff_f)
 else if(NS .and. lon_diff_f.gt.0.)then
          gamma=atan(lat_diff_f/lon_diff_f)+3.14/2.
      else if(NS .and. lon_diff_f.eq.0)then
             gamma=0.
           end if
      end if
 end if

;;;;Make use of vector algebra and use a rotation matrix :-);;;;;;;;;
 print(gamma*180/3.14)
 wwg_along      = cos(gamma)*uga(:,:,1:dimP(2)-2)+sin(gamma)*vga(:,1:dimP(1)-2,:)
 wwg_across     = -sin(gamma)*uga(:,:,1:dimP(2)-2)+cos(gamma)*vga(:,1:dimP(1)-2,:)
 wwg_along@description  = "Along geostrophic Wind"
 wwg_across@description = "Across geostrophic Wind"


 u             = wrf_user_getvar(a,"U",it)      ; 3D U at mass points
 v             = wrf_user_getvar(a,"V",it)      ; 3D V at mass points
 theta         = wrf_user_getvar(a,"theta",it)

;;;;interpolate it to masspoints myself, since i don't know what wrf_user_getvar is really doing;;;;;;; 

 ua=(u(1:dimP(0)-2,1:dimP(1)-2,0:dimP(2)-1)+u(1:dimP(0)-2,1:dimP(1)-2,1:dimP(2)))/2
 va=(v(1:dimP(0)-2,0:dimP(1)-1,1:dimP(2)-2)+v(1:dimP(0)-2,1:dimP(1),1:dimP(2)-2))/2


 ww_along      = cos(gamma)*ua(:,:,1:dimP(2)-2)+sin(gamma)*va(:,1:dimP(1)-2,:)
 ww_across     = -sin(gamma)*ua(:,:,1:dimP(2)-2)+cos(gamma)*va(:,1:dimP(1)-2,:)
;calculating the ageostrophic components along and across the cross section


wwageo_along  = ww_along-wwg_along
wwageo_across = ww_across-wwg_across
wwageo_along@description = "Along ageostrophic wind"
wwageo_across@description= "Across ageostrophic wind" 




all_winds=new((/dimP(0)-2,dimP(1)-2,dimP(2)-2,6/),typeof(ww_along))
all_winds(:,:,:,0)=ww_along

all_winds(:,:,:,0)=ww_along
all_winds(:,:,:,1)=ww_across
all_winds(:,:,:,2)=wwg_along
all_winds(:,:,:,3)=wwg_across
all_winds(:,:,:,4)=wwageo_along
all_winds(:,:,:,5)=wwageo_across


return(all_winds)

end