# Metview Macro

# **************************** LICENSE START ***********************************
#
# Copyright 2012 ECMWF. This software is distributed under the terms
# of the Apache License version 2.0. In applying this license, ECMWF does not
# waive the privileges and immunities granted to it by virtue of its status as
# an Intergovernmental Organization or submit itself to any jurisdiction.
#
# ***************************** LICENSE END ************************************

# **************************************************************************
# Function      : mvl_geopotential_on_ml
#
# Syntax        : fieldset mvl_geopotential_on_ml (T:fieldset, q:fieldset, lnsp:fieldset, zs:fieldset)
#
# Author (date) : Iain Russell (12/07/2011)
#
# Category      : COMPUTATION
#
# OneLineDesc   : Computes geopotential on model levels
#
# Description   : Computes geopotential on model levels.
#                 Based on code from Nils Wedi, the IFS documentation:
#                 http://www.ecmwf.int/research/ifsdocs/DYNAMICS/Chap2_Discretization4.html
#                 and an optimised implementation by Dominique Lucas.
#
# Parameters    : T    - fieldset of Temperature on model levels in ascending numeric order (e.g. 1-91)
#                 q    - fieldset of specific humidity on model levels in ascending numeric order (e.g. 1-91)
#                 lnsp - field of log of surface pressure on model level 1
#                 zs   - field of geopotential on model level 1 (available from MARS)
#                 All fields must be GRIDDED data - no spherical harmonics, and they must all be
#                 on the same grid, with the same number of points.
#
# Return Value  : A fieldset of geopotential on model levels
#
# Dependencies  : None
#
# Example Usage : 
#       r = (date: -1, time: 12, levtype: "ml", grid: [1.5,1.5])
#       T    = retrieve(r,levelist: [1,"to",91],param: "t")
#       q    = retrieve(r,levelist: [1,"to",91],param: "q")
#       zs   = retrieve(r,levelist: 1,param: "z")
#       lnsp = retrieve(r,levelist: 1,param: "lnsp")
#       z_ml = mvl_geopotential_on_ml(T, q, lnsp, zs)
#
# **************************************************************************

function mvl_geopotential_on_ml (T_fs:fieldset, q_fs:fieldset, lnsp_fs:fieldset, zs_fs:fieldset)

    Rd = 287.06
    g  = 9.80665


    # do some basic checking on the input data - only check the first field of each fielset for efficiency

    gridtype_t = grib_get_string(T_fs[1],    'gridType')
    gridtype_q = grib_get_string(q_fs[1],    'gridType')
    gridtype_l = grib_get_string(lnsp_fs[1], 'gridType')
    gridtype_z = grib_get_string(zs_fs[1],   'gridType')

    if (gridtype_t = 'sh' or gridtype_q = 'sh' or gridtype_l = 'sh' or gridtype_z = 'sh') then
        fail ('mvl_geopotential_on_ml: fields must be gridded, not spectral')
    end if



    levelSize = count(T_fs)    # how many levels are we computing?

    zs   = values(zs_fs)     # extract the data for this field as a vector
    lnsp = values(lnsp_fs)   # extract the data for this field as a vector


    pv = grib_get_double_array(zs_fs, 'pv') # for computing pressures at model levels
    sp = exp(lnsp)                          # surface pressure
    z_h = zs                                # orography


    # get the coefficients for computing the pressures
    
    pv = grib_get_double_array(zs_fs, 'pv')
    A = pv[1, levelSize+1]                #e.g. L-91: A=1..92, B=93..184
    B = pv[levelSize+2, levelSize+levelSize+2]


    z_f = z_h        # output vector
    z_out_fs = nil   # output fieldset - the end result



    # compute the bottom pressure (on half-levels)

    Ph_levplusone = (A[levelSize+1] + (B[levelSize+1] * sp)) # initialise to the lowest numbered level


    # We want to integrate up into the atmosphere, starting at the ground
    # so we start at the lowest level (highest number) and keep
    # accumulating the height as we go.
    # See the IFS documentation:
    # http://www.ecmwf.int/research/ifsdocs/DYNAMICS/Chap2_Discretization4.html
    # For speed and file I/O, we perform the computations with vectors instead
    # of fieldsets.



    for lev = levelSize to 1 by -1 do

        # we assume that the data are in (numerically) ascending level order - check if true

        T_level = grib_get_long(T_fs[lev], 'level')
        q_level = grib_get_long(q_fs[lev], 'level')

        if (T_level <> lev) then
            fail ('T field index ', lev, ' should be level ', lev, ' but it is ', T_level)
        else if (q_level <> lev) then
            fail ('q field index ', lev, ' should be level ', lev, ' but it is ', T_level)
        end if


        q = values(q_fs[lev]) # extract the data for this field as a vector
        T = values(T_fs[lev]) # extract the data for this field as a vector


        # quick check for data consistency

        if ((count(q) <> count(T)) or
            (count(q) <> count(sp)) or
            (count(q) <> count(z_h))) then
            errmsg = 'mvl_geopotential_on_ml: T, Q, LNSP and Z, must have the same number of grid points (they have ' &
                      count(T) & ', ' & count(q) & ', ' & count(sp) & ' and ' & count(z_h) & ' respectively).'
            fail(errmsg)
        end if


        T = T*(1.+0.609133*q)   # compute moist temperature
        q = 0                   # free the memory for 'q'


        # compute the pressures (on half-levels)

        Ph_lev = (A[lev] + (B[lev] * sp))


        if lev = 1 then
            dlogP = log(Ph_levplusone/0.1) # Ph[lev] is zero, so don't divide by it
            alpha = log(2)
        else
            dlogP = log(Ph_levplusone/Ph_lev)
            dP    = Ph_levplusone-Ph_lev
            alpha = 1 - Ph_lev/dP*dlogP
        end if

        TRd = T*Rd
        T   = 0      # free memory for 'T'


        # z_f is the geopotential of this full level
        # integrate from previous (lower) half-level z_h to the full level

        z_f = z_h + (TRd*alpha)


        # z_h is the geopotential of 'half-levels'
        # integrate z_h to next half level

        z_h = z_h + (TRd*dlogP)


        # store the result (z_f) in a field and add to the output fieldset
        # (add it to the front, not the end, because we are going 'backwards'
        # through the fields)

        out_field = set_values(T_fs[lev], z_f)
        out_field = grib_set_long(out_field, ['paramId', 129,
                                              'generatingProcessIdentifier',128,
                                              'level', lev])
        z_out_fs = out_field & z_out_fs

        Ph_levplusone  = Ph_lev    # store for the next iteration, will be equivalent to (A[lev+1] + (B[lev+1] * sp))

    end for

    return z_out_fs

end mvl_geopotential_on_ml
