NumRu::GAnalysis::QG_common

# File ../../lib/numru/ganalysis/qg.rb, line 154
def cut_bottom(z)
  pdim = Met.find_prs_d(z)
  plev = z.coord(pdim).val
  if plev[0] - plev[1] > 0
    z[ *([true]*pdim + [1..-1,false]) ]
  else
    z[ *([true]*pdim + [0..-2,false]) ]
  end
end

# File ../../lib/numru/ganalysis/qg.rb, line 127
def extend_bottom(z, val_extended=nil)
  pdim = Met.find_prs_d(z)
  plev = z.coord(pdim).val
  raise("Only one pressure level is found; 2 or more needed") if (plev==1)
  bottom_first = ( plev[0] - plev[1] > 0 )
  np = z.shape[pdim]
  idx = (0...np).collect{|i| i}
  if bottom_first       # The first level is the bottom one
    idx.unshift(0)      # idx => [0,0,1,2,...,np-1]
    ihb = 0             # index of the extended bottom level
    dp = plev[0] - plev[1]
    phb = plev[0] + dp  # pressure of the extended bottom level
  else                # The last level is the bottom one
    idx.push(np-1)     # idx => [0,1,2,...,np-1,np-1]
    ihb = np            # index of the extended bottom level
    dp = plev[-1] - plev[-2]
    phb = plev[-1] + dp  # pressure of the extended bottom level
  end
  ze = z[ *([true]*pdim + [idx,false]) ].copy  # add one level below
  ze.coord(pdim)[ihb] = phb
  if val_extended
    ze[ *([true]*pdim + [ihb,false]) ] = val_extended
  end
  ze
end

# File ../../lib/numru/ganalysis/qg.rb, line 24
def gp2gpref(gp)
  gpref = Planet::ave_s(gp)  # horizontal ave (spherical)
  if gpref.rank >= 2
    # likely a time sequence. need to reduce more.
    pdim = Met.find_prs_d(gpref)
    idxs = (0...gpref.rank).collect{|i| i}
    idxs.delete(pdim)
    gpref = gpref.mean(*idxs)
  end
  gpref.name = "gpref"
  gpref.long_name = "Reference geopotential"
  gpref
end

# File ../../lib/numru/ganalysis/qg.rb, line 65
def gpd2qzz(gp, b)
  bunits = Units["s-2"]   # this is assumed!
  pdim = Met.find_prs_d(gp)
  p = gp.coord(pdim)
  z = LogP.p2z(p)
  zunits = z.units
  g = Derivative::b_expand_linear_ext( gp.val, pdim )
  z = Derivative::b_expand_linear_ext( z.val, 0 )
  p = Derivative::b_expand_linear_ext( p.val, 0 )
  b = b.val
  b = b.to_na if b.respond_to?(:to_na)  # likely a NArrayMiss
  b = Derivative::b_expand_linear_ext( b, 0 )
  pb = p/b

  pbm = (pb[0..-2] + pb[1..-1]) / 2.0  # pb_{i+1/2} (for i=0..-2)
  pbm01 = pbm[0..-2]                   # pb_{i-1/2} (for i=1..-2)
  pbm12 = pbm[1..-1]                   # pb_{i+1/2} (for i=1..-2)
  dz20 = z[2..-1] - z[0..-3]           # z_{i+1} - z_{i-1} (for i=1..-2)
  dz21 = z[2..-1] - z[1..-2]           # z_{i+1} - z_{i}   (for i=1..-2)
  dz10 = z[1..-2] - z[0..-3]           # z_{i} - x_{i-1}   (for i=1..-2)
  pc = p[1..-2]                        # p_{i}   (for i=1..-2)

  a2 = 2*pbm12/(dz21*dz20)/pc
  a0 = 2*pbm01/(dz10*dz20)/pc
  a1 = -a2 - a0

  to_rankD = [1]*pdim + [true] + [1]*(gp.rank-pdim-1)
  a2 = a2.reshape(*to_rankD)
  a1 = a1.reshape(*to_rankD)
  a0 = a0.reshape(*to_rankD)

  vqzz =   g[ *([true]*pdim+[2..-1,false]) ] * a2                 + g[ *([true]*pdim+[1..-2,false]) ] * a1                 + g[ *([true]*pdim+[0..-3,false]) ] * a0

  qzz = gp.copy
  qzz.data.replace_val(vqzz)
  qzz.name = "qzz"
  qzz.long_name = "z-deriv term in QG PV"
  qzz.units = qzz.units / zunits**2 / bunits
  qzz
end

# File ../../lib/numru/ganalysis/qg.rb, line 40
def gph2gpd_gpref(gph)
  gp = gph * Met::g
  gpref = gp2gpref(gp)
  gpd = gp - gpref
  gpd.name = "gpd"
  gpd.long_name = "Geopotential deviation"
  [gpd, gpref]
end

# File ../../lib/numru/ganalysis/qg.rb, line 18
def gph2gpref(gph)
  gp2gpref(gph) * Met::g
end

# File ../../lib/numru/ganalysis/qg.rb, line 51
def gpref2n2(gpref)
  gp_z = LogP.pcdata_dz( gpref )
  gp_zz = LogP.pcdata_dz2( gpref )
  gp_zz[0] = gp_zz[1]   # At boundary, it's safer to extend lapse rate
  gp_zz[-1] = gp_zz[-2] # At boundary, it's safer to extend lapse rate
  n2 = gp_zz + gp_z * (Met::Kappa / LogP.h)
  n2.name = "N2"
  n2.long_name = "N**2 (log-p)"
  #p "@@@@@ N2 @@@@",n2.coord(0).val.to_a, n2.val.sqrt.to_a
  n2
end

# File ../../lib/numru/ganalysis/qg.rb, line 185
def div_waf(fx, fy, fz, bottom_treatment=true)
  cosphi = cos_phi(fx)
  p = Met.get_prs(fx)

  fz_z = LogP.pcdata_dz( fz*p ) / p

  #>>>>>> the lowest layer treatment consistent with qb, in which
  # geopotential (or stream function) is extended by extend_bottom. 
  # Assumption: the first level is the lowest (bottom) one
  if bottom_treatment
    # using the relation p^{-1} d/dz = -H^{-1} d/dp 
    # and assuming fz=0 below the bottom (the "underground" level),
    # p^{-1} d/dz (p fz) = -H^{-1} d/dp (p fz),
    # which is H^{-1} p fz / delta_p, at the lowest level with a
    # "thickness" of delta_p. 
    w = p[0..1].val
    dp = w[1] - w[0]
    p0 = w[0]
    pdim = Met.find_prs_d(fz)
    sel0 = [true]*pdim + [0,false]  # to specify the first level
    fz_z[*sel0] = fz[*sel0]*p0 / (LogP.h*dp)
  end
  #<<<<<<

  divh = ( div_h(fx, fy) + fz_z ) / cosphi
         # ^ div_h is defined in QG, QG_sphere,..., but not in QG_common
  divh.name = "divwaf"
  divh.long_name = "div of waf (#{fx.name},..)"
  divh
end