Ejemplo n.º 1
0
def interp_WWIND_UVWIND_py(DWIND, DWIND_km1, WWIND, WWIND_dm1, WWIND_pm1,
                           WWIND_pp1, WWIND_pm1_dm1, WWIND_pp1_dm1, COLP_NEW,
                           COLP_NEW_dm1, COLP_NEW_pm1, COLP_NEW_pp1,
                           COLP_NEW_pm1_dm1, COLP_NEW_pp1_dm1, A, A_dm1, A_pm1,
                           A_pp1, A_pm1_dm1, A_pp1_dm1, dsigma, dsigma_km1,
                           rigid_wall, p_ind, np, k):
    """
    Interpolate WWIND * UVWIND (U or V, depending on direction) 
    onto position of repective horizontal wind.
    d inds (e.g. dm1 = d minus 1) are in direction of hor. wind
    vector.
    p inds are perpendicular to direction of hor. wind vector.
    if rigid_wall: Special BC for rigid wall parallel to flow
    direction according to hint from Mark Jacobson during mail
    conversation. np is number of grid cells and p_ind current
    index in direction perpeindular to flow.
    """
    if k == wp_int(0) or k == nzs - wp_int(1):
        WWIND_DWIND = wp(0.)
    else:
        # left rigid wall
        if rigid_wall and (p_ind == nb):
            COLPAWWIND_ds_ks = wp(0.25) * (
                COLP_NEW_pp1_dm1 * A_pp1_dm1 * WWIND_pp1_dm1 +
                COLP_NEW_pp1 * A_pp1 * WWIND_pp1 +
                COLP_NEW_dm1 * A_dm1 * WWIND_dm1 + COLP_NEW * A * WWIND)
        # right rigid wall
        elif rigid_wall and (p_ind == np):
            COLPAWWIND_ds_ks = wp(0.25) * (
                COLP_NEW_dm1 * A_dm1 * WWIND_dm1 + COLP_NEW * A * WWIND +
                COLP_NEW_pm1_dm1 * A_pm1_dm1 * WWIND_pm1_dm1 +
                COLP_NEW_pm1 * A_pm1 * WWIND_pm1)
        # inside domain (not at boundary of perpendicular dimension)
        else:
            COLPAWWIND_ds_ks = wp(0.125) * (
                COLP_NEW_pp1_dm1 * A_pp1_dm1 * WWIND_pp1_dm1 + COLP_NEW_pp1 *
                A_pp1 * WWIND_pp1 + wp(2.) * COLP_NEW_dm1 * A_dm1 * WWIND_dm1 +
                wp(2.) * COLP_NEW * A * WWIND + COLP_NEW_pm1_dm1 * A_pm1_dm1 *
                WWIND_pm1_dm1 + COLP_NEW_pm1 * A_pm1 * WWIND_pm1)

        # interpolate hor. wind on vertical interface
        DWIND_ks = ((dsigma * DWIND_km1 + dsigma_km1 * DWIND) /
                    (dsigma + dsigma_km1))

        # combine
        WWIND_DWIND = COLPAWWIND_ds_ks * DWIND_ks

    return (WWIND_DWIND)
Ejemplo n.º 2
0
def turb_flux_tendency_py(PHI, PHI_kp1, PHI_km1, PHIVB, PHIVB_kp1, VAR,
                          VAR_kp1, VAR_km1, KVAR, KVAR_kp1, RHO, RHOVB,
                          RHOVB_kp1, COLP, surf_flux_VAR, k):
    """
    Vertical turbulent transport. 
    """
    ALT = PHI / con_g
    ALT_kp1 = PHI_kp1 / con_g
    ALT_km1 = PHI_km1 / con_g
    ALTVB = PHIVB / con_g
    ALTVB_kp1 = PHIVB_kp1 / con_g

    if k == wp_int(0):
        #dVARdt_TURB = wp(0.)
        dVARdt_TURB = COLP * ((+wp(0.) -
                               ((VAR - VAR_kp1) /
                                (ALT - ALT_kp1) * RHOVB_kp1 * KVAR_kp1)) /
                              ((ALTVB - ALTVB_kp1) * RHO))
    elif k == nz - 1:
        #dVARdt_TURB = wp(0.)
        dVARdt_TURB = COLP * (
            (+((VAR_km1 - VAR) /
               (ALT_km1 - ALT) * RHOVB * KVAR) + surf_flux_VAR) /
            ((ALTVB - ALTVB_kp1) * RHO))
    else:
        #dVARdt_TURB = wp(0.)
        dVARdt_TURB = COLP * ((+((VAR_km1 - VAR) /
                                 (ALT_km1 - ALT) * RHOVB * KVAR) -
                               ((VAR - VAR_kp1) /
                                (ALT - ALT_kp1) * RHOVB_kp1 * KVAR_kp1)) /
                              ((ALTVB - ALTVB_kp1) * RHO))
    return (dVARdt_TURB)
Ejemplo n.º 3
0
def launch_numba_cpu(A, dsigma, POTT_dif_coef, dPOTTdt, POTT, UFLX, VFLX, COLP,
                     POTTVB, WWIND, COLP_NEW, PHI, PHIVB, KHEAT, RHO, RHOVB,
                     SSHFLX, dPOTTdt_TURB, dPOTTdt_RAD):

    for i in prange(nb, nx + nb):
        for j in range(nb, ny + nb):
            for k in range(wp_int(0), nz):
                dPOTTdt[i  ,j  ,k], dPOTTdt_TURB[i  ,j  ,k] = \
                    add_up_tendencies(POTT[i  ,j  ,k],
                        POTT        [i-1,j  ,k  ], POTT         [i+1,j  ,k  ],
                        POTT        [i  ,j-1,k  ], POTT         [i  ,j+1,k  ],
                        POTT        [i  ,j  ,k-1], POTT         [i  ,j  ,k+1],
                        UFLX        [i  ,j  ,k  ], UFLX[i+1,j  ,k],
                        VFLX        [i  ,j  ,k  ], VFLX[i  ,j+1,k],
                        COLP        [i  ,j  ,0  ],
                        COLP        [i-1,j  ,0  ], COLP[i+1,j  ,0],
                        COLP        [i  ,j-1,0  ], COLP[i  ,j+1,0],
                        POTTVB      [i  ,j  ,k  ], POTTVB[i  ,j  ,k+1],
                        WWIND       [i  ,j  ,k  ], WWIND[i  ,j  ,k+1],
                        COLP_NEW    [i  ,j  ,0  ],

                        PHI         [i  ,j  ,k  ], PHI        [i  ,j  ,k+1],
                        PHI         [i  ,j  ,k-1], PHIVB      [i  ,j  ,k  ],
                        PHIVB       [i  ,j  ,k+1],
                        KHEAT       [i  ,j  ,k  ], KHEAT      [i  ,j  ,k+1],
                        RHO         [i  ,j  ,k  ], RHOVB      [i  ,j  ,k  ],
                        RHOVB       [i  ,j  ,k+1], SSHFLX     [i  ,j  ,0  ],

                        dPOTTdt_RAD [i  ,j  ,k],
                        A           [i  ,j  ,0  ],
                        dsigma      [0  ,0  ,k  ], POTT_dif_coef[0  ,0  ,k],
                        k)
Ejemplo n.º 4
0
def vert_adv_py(VARVB, VARVB_kp1, WWIND, WWIND_kp1, COLP_NEW, dsigma, k):
    """
    Vertical advection. 
    """

    if k == wp_int(0):
        return (COLP_NEW * (-WWIND_kp1 * VARVB_kp1) / dsigma)
    elif k == nz:
        return (COLP_NEW * (+WWIND * VARVB) / dsigma)
    else:
        return (COLP_NEW * (+WWIND * VARVB - WWIND_kp1 * VARVB_kp1) / dsigma)
Ejemplo n.º 5
0
def launch_numba_cpu(KMOM, KHEAT, PHIVB, HSURF, PHI, QV, WINDX, WINDY, POTTVB,
                     POTT):

    for i in prange(0, nx + 2 * nb):
        for j in range(0, ny + 2 * nb):
            for k in range(wp_int(1), nzs - 1):
                KMOM[i, j, k], KHEAT[i, j, k] = run_all(
                    PHIVB[i, j, k], HSURF[i, j, 0], PHI[i, j, k - 1],
                    PHI[i, j, k], QV[i, j, k - 1], QV[i, j, k], WINDX[i, j,
                                                                      k - 1],
                    WINDX[i, j, k], WINDY[i, j, k - 1], WINDY[i, j, k],
                    POTTVB[i, j, k], POTT[i, j, k - 1], POTT[i, j, k], k)
Ejemplo n.º 6
0
def diag_secondary_cpu(POTTVB, TAIRVB, PVTFVB, 
                        COLP, PAIR, PAIRVB, PHI, POTT, 
                        TAIR, RHO, RHOVB, PVTF,
                        UWIND, VWIND, WINDX, WINDY, WIND):
    for i in prange(0,nx+2*nb):
        for j in range(0,ny+2*nb):
            for k in range(wp_int(0),nzs):
                PAIRVB[i,j,k] = wp(100000.)*(
                                PVTFVB[i,j,k])**(wp(1.)/con_kappa)
                TAIRVB[i,j,k] = POTTVB[i,j,k] * PVTFVB[i,j,k]
                RHOVB [i,j,k] = PAIRVB[i,j,k] / (con_Rd * TAIRVB[i,j,k])

            for k in range(wp_int(0),nz):
                TAIR [i,j,k] = POTT[i,j,k] * PVTF[i,j,k]
                PAIR [i,j,k] = wp(100000.)*(PVTF[i,j,k])**(wp(1.)/con_kappa)
                RHO  [i,j,k] = PAIR[i,j,k] / (con_Rd * TAIR[i,j,k])
                WINDX[i,j,k] = (UWIND[i  ,j  ,k] + UWIND[i+1,j  ,k])/wp(2.)
                WINDY[i,j,k] = (VWIND[i  ,j  ,k] + VWIND[i  ,j+1,k])/wp(2.) 
                WIND [i,j,k] = (
                    ((UWIND[i  ,j  ,k] + UWIND[i+1,j  ,k])/wp(2.))**wp(2.) +
                    ((VWIND[i  ,j  ,k] + VWIND[i  ,j+1,k])/wp(2.))**wp(2.) 
                              ) ** (wp(1.)/wp(2.))
Ejemplo n.º 7
0
def launch_numba_cpu(UFLX, VFLX, FLXDIV, UWIND, VWIND, WWIND, COLP, dCOLPdt,
                     COLP_NEW, COLP_OLD, dyis, dxjs, dsigma, sigma_vb, A, dt):

    for i in prange(nb, nx + nb):
        for j in range(nb, ny + nb):
            for k in range(wp_int(0), nz):
                # MOMENTUM FLUXES
                ###############################################################
                UFLX_i = calc_UFLX(UWIND[i, j, k], COLP[i, j, 0],
                                   COLP[i - 1, j, 0], dyis[i, j, 0])
                UFLX_ip1 = calc_UFLX(UWIND[i + 1, j, k], COLP[i + 1, j, 0],
                                     COLP[i, j, 0], dyis[i + 1, j, 0])

                VFLX_j = calc_VFLX(VWIND[i, j, k], COLP[i, j, 0],
                                   COLP[i, j - 1, 0], dxjs[i, j, 0])
                VFLX_jp1 = calc_VFLX(VWIND[i, j + 1, k], COLP[i, j + 1, 0],
                                     COLP[i, j, 0], dxjs[i, j + 1, 0])

                UFLX[i, j, k] = UFLX_i
                VFLX[i, j, k] = VFLX_j

                ## MOMENTUM FLUX DIVERGENCE
                ###############################################################
                FLXDIV[i, j, k] = calc_FLXDIV(UFLX_i, UFLX_ip1, VFLX_j,
                                              VFLX_jp1, dsigma[0, 0,
                                                               k], A[i, j, 0])

    ## COLUMN PRESSURE TENDENCY
    ############################################################################
    if i_COLP_main_switch:
        dCOLPdt[:, :, 0] = -FLXDIV.sum(axis=2)
    else:
        dCOLPdt[:, :, 0] = wp(0.)

    ### PRESSURE TIME STEP
    ############################################################################
    COLP_NEW[:, :, 0] = COLP_OLD[:, :, 0] + dt * dCOLPdt[:, :, 0]

    ### VERTICAL WIND
    ############################################################################
    for i in prange(nb, nx + nb):
        for j in range(nb, ny + nb):
            flxdivsum = FLXDIV[i, j, 0]
            for k in prange(1, nz):
                WWIND[i, j, k] = (
                    -flxdivsum / COLP_NEW[i, j, 0] -
                    sigma_vb[0, 0, k] * dCOLPdt[i, j, 0] / COLP_NEW[i, j, 0])
                flxdivsum += FLXDIV[i, j, k]
Ejemplo n.º 8
0
def diag_POTTVB_cpu(POTTVB, POTT, PVTF, PVTFVB):
    for i in prange(0,nx+2*nb):
        for j in range(0,ny+2*nb):
            for k in range(wp_int(1),nzs-1):
                POTTVB[i,j,k] =   (
                        +   (PVTFVB[i,j,k] - PVTF  [i,j,k-1]) * POTT[i,j,k-1]
                        +   (PVTF  [i,j,k] - PVTFVB[i,j,k  ]) * POTT[i,j,k  ]
                                ) / (PVTF[i,j,k] - PVTF[i,j,k-1])
                if k == 1:
                    # extrapolate model top POTTVB
                    POTTVB[i,j,k-1] = POTT[i,j,k-1] - ( 
                                            POTTVB[i,j,k] - POTT[i,j,k-1] )
                elif k == nzs-2:
                    # extrapolate model bottom POTTVB
                    POTTVB[i,j,k+1] = POTT[i,j,k  ] - (
                                            POTTVB[i,j,k] - POTT[i,j,k  ] )
Ejemplo n.º 9
0
def diag_PVTF_cpu(COLP, PVTF, PVTFVB, sigma_vb):

    for i in prange(0,nx+2*nb):
        for j in range(0,ny+2*nb):
            for k in range(wp_int(0),nz):
        
                pairvb_km12 = pair_top + sigma_vb[0,0,k  ] * COLP[i,j,0]
                pairvb_kp12 = pair_top + sigma_vb[0,0,k+1] * COLP[i,j,0]
                
                PVTF[i,j,k] = wp(1.)/(wp(1.)+con_kappa) * (
                        pow( pairvb_kp12/wp(100000.) , con_kappa ) * pairvb_kp12 -
                        pow( pairvb_km12/wp(100000.) , con_kappa ) * pairvb_km12
                                ) /( pairvb_kp12 - pairvb_km12 )

                PVTFVB[i,j,k] = pow( pairvb_km12/wp(100000.) , con_kappa )
                if k == nz-1:
                    PVTFVB[i,j,k+1] = pow( pairvb_kp12/wp(100000.) , con_kappa )
Ejemplo n.º 10
0
def launch_numba_cpu(A, dsigma, moist_dif_coef,
                    dQVdt, dQVdt_TURB, QV, dQCdt, QC, UFLX, VFLX, COLP,
                    WWIND, COLP_NEW,
                    PHI, PHIVB, KHEAT, RHO, RHOVB, SLHFLX):

    for i in prange(nb,nx+nb):
        for j in range(nb,ny+nb):
            for k in range(wp_int(0),nz):
                dQVdt[i  ,j  ,k], dQVdt_TURB[i  ,j  ,k], dQCdt[i  ,j  ,k] = \
                    add_up_tendencies(
                        QV          [i  ,j  ,k  ],
                        QV          [i-1,j  ,k  ], QV       [i+1,j  ,k  ],
                        QV          [i  ,j-1,k  ], QV       [i  ,j+1,k  ],
                        QV          [i  ,j  ,k-1], QV       [i  ,j  ,k+1],
                        QC          [i  ,j  ,k  ],
                        QC          [i-1,j  ,k  ], QC       [i+1,j  ,k  ],
                        QC          [i  ,j-1,k  ], QC       [i  ,j+1,k  ],
                        QC          [i  ,j  ,k-1], QC       [i  ,j  ,k+1],
                        UFLX        [i  ,j  ,k  ], UFLX     [i+1,j  ,k  ],
                        VFLX        [i  ,j  ,k  ], VFLX     [i  ,j+1,k  ],
                        COLP        [i  ,j  ,0  ],
                        COLP        [i-1,j  ,0  ], COLP     [i+1,j  ,0  ],
                        COLP        [i  ,j-1,0  ], COLP     [i  ,j+1,0  ],
                        WWIND       [i  ,j  ,k  ], WWIND    [i  ,j  ,k+1],
                        COLP_NEW    [i  ,j  ,0  ],

                        PHI         [i  ,j  ,k  ], PHI      [i  ,j  ,k+1],
                        PHI         [i  ,j  ,k-1], PHIVB    [i  ,j  ,k  ],
                        PHIVB       [i  ,j  ,k+1], 
                        KHEAT       [i  ,j  ,k  ], KHEAT    [i  ,j  ,k+1],
                        RHO         [i  ,j  ,k  ], RHOVB    [i  ,j  ,k  ],
                        RHOVB       [i  ,j  ,k+1], SLHFLX   [i  ,j  ,0  ],

                        A           [i  ,j  ,0  ],
                        dsigma      [0  ,0  ,k  ], moist_dif_coef[0  ,0  ,k],
                        k)
Ejemplo n.º 11
0
def interp_KMOM_dUVWINDdz_py(
        DWIND, DWIND_km1, KMOM, KMOM_dm1, KMOM_pm1, KMOM_pp1, KMOM_pm1_dm1,
        KMOM_pp1_dm1, RHOVB, RHOVB_dm1, RHOVB_pm1, RHOVB_pp1, RHOVB_pm1_dm1,
        RHOVB_pp1_dm1, PHI, PHI_dm1, PHI_pm1, PHI_pp1, PHI_pm1_dm1,
        PHI_pp1_dm1, PHI_km1, PHI_km1_dm1, PHI_km1_pm1, PHI_km1_pp1,
        PHI_km1_pm1_dm1, PHI_km1_pp1_dm1, COLP_NEW, COLP_NEW_dm1, COLP_NEW_pm1,
        COLP_NEW_pp1, COLP_NEW_pm1_dm1, COLP_NEW_pp1_dm1, A, A_dm1, A_pm1,
        A_pp1, A_pm1_dm1, A_pp1_dm1, dsigma, dsigma_km1, rigid_wall, p_ind, np,
        k):
    """
    Interpolate KMOM * UVWIND (U or V, depending on direction) 
    onto position of repective horizontal wind.
    d inds (e.g. dm1 = d minus 1) are in direction of hor. wind
    vector.
    p inds are perpendicular to direction of hor. wind vector.
    if rigid_wall: Special BC for rigid wall parallel to flow
    direction according to hint from Mark Jacobson during mail
    conversation. np is number of grid cells and p_ind current
    index in direction perpeindular to flow.
    """
    if k == wp_int(0) or k == nzs - wp_int(1):
        KMOM_DWIND = wp(0.)
    else:
        # left rigid wall
        if rigid_wall and (p_ind == nb):
            COLPAKMOM_ds_ks = wp(0.25) * (
                RHOVB_pp1_dm1 * COLP_NEW_pp1_dm1 * A_pp1_dm1 * KMOM_pp1_dm1 +
                RHOVB_pp1 * COLP_NEW_pp1 * A_pp1 * KMOM_pp1 + RHOVB_dm1 *
                COLP_NEW_dm1 * A_dm1 * KMOM_dm1 + RHOVB * COLP_NEW * A * KMOM)
            ALT_ds_km1 = wp(0.25) * (PHI_km1_pp1_dm1 + PHI_km1_pp1 +
                                     PHI_km1_dm1 + PHI_km1) / con_g
            ALT_ds = wp(0.25) * (PHI_pp1_dm1 + PHI_pp1 + PHI_dm1 + PHI) / con_g
        # right rigid wall
        elif rigid_wall and (p_ind == np):
            COLPAKMOM_ds_ks = wp(0.25) * (
                RHOVB_dm1 * COLP_NEW_dm1 * A_dm1 * KMOM_dm1 +
                RHOVB * COLP_NEW * A * KMOM +
                RHOVB_pm1_dm1 * COLP_NEW_pm1_dm1 * A_pm1_dm1 * KMOM_pm1_dm1 +
                RHOVB_pm1 * COLP_NEW_pm1 * A_pm1 * KMOM_pm1)
            ALT_ds_km1 = wp(0.25) * (PHI_km1_dm1 + PHI_km1 + PHI_km1_pm1_dm1 +
                                     PHI_km1_pm1) / con_g
            ALT_ds = wp(0.25) * (PHI_dm1 + PHI + PHI_pm1_dm1 + PHI_pm1) / con_g
        # inside domain (not at boundary of perpendicular dimension)
        else:
            COLPAKMOM_ds_ks = wp(0.125) * (
                RHOVB_pp1_dm1 * COLP_NEW_pp1_dm1 * A_pp1_dm1 * KMOM_pp1_dm1 +
                RHOVB_pp1 * COLP_NEW_pp1 * A_pp1 * KMOM_pp1 +
                wp(2.) * RHOVB_dm1 * COLP_NEW_dm1 * A_dm1 * KMOM_dm1 +
                wp(2.) * RHOVB * COLP_NEW * A * KMOM +
                RHOVB_pm1_dm1 * COLP_NEW_pm1_dm1 * A_pm1_dm1 * KMOM_pm1_dm1 +
                RHOVB_pm1 * COLP_NEW_pm1 * A_pm1 * KMOM_pm1)
            ALT_ds_km1 = wp(0.125) * (PHI_km1_pp1_dm1 + PHI_km1_pp1 +
                                      wp(2.) * PHI_km1_dm1 + wp(2.) * PHI_km1 +
                                      PHI_km1_pm1_dm1 + PHI_km1_pm1) / con_g
            ALT_ds = wp(0.125) * (PHI_pp1_dm1 + PHI_pp1 + wp(2.) * PHI_dm1 +
                                  wp(2.) * PHI + PHI_pm1_dm1 + PHI_pm1) / con_g

        # interpolate hor. wind on vertical interface
        dDWINDdz_ks = ((DWIND_km1 - DWIND) / (ALT_ds_km1 - ALT_ds))

        # combine
        KMOM_dDWINDdz = COLPAKMOM_ds_ks * dDWINDdz_ks

    return (KMOM_dDWINDdz)
Ejemplo n.º 12
0
def launch_numba_cpu_main(dUFLXdt, UFLX, UWIND, VWIND, BFLX_3D, CFLX_3D,
                          DFLX_3D, EFLX_3D, PHI, PHIVB, COLP, POTT, PVTF,
                          PVTFVB, WWIND_UWIND, KMOM_dUWINDdz, RHO,
                          dUFLXdt_TURB, SMOMXFLX, corf_is, lat_is_rad,
                          dlon_rad, dlat_rad, dyis, dsigma, sigma_vb,
                          UVFLX_dif_coef):

    for i in prange(nb, nxs + nb):
        for j in range(nb, ny + nb):
            for k in range(wp_int(0), nz):
                # Prepare momentum fluxes and set boundary conditions if
                # necessary
                BFLX = BFLX_3D[i, j, k]
                CFLX = CFLX_3D[i, j, k]
                EFLX = EFLX_3D[i, j, k]
                DFLX_jp1 = DFLX_3D[i, j + 1, k]
                CFLX_jp1 = CFLX_3D[i, j + 1, k]

                # BCx i
                if i == nb:
                    BFLX_im1 = BFLX_3D[nx, j, k]
                    DFLX_im1 = DFLX_3D[nx, j, k]
                    EFLX_im1_jp1 = EFLX_3D[nx, j + 1, k]
                else:
                    BFLX_im1 = BFLX_3D[i - 1, j, k]
                    DFLX_im1 = DFLX_3D[i - 1, j, k]
                    EFLX_im1_jp1 = EFLX_3D[i - 1, j + 1, k]

                # BCy js
                if j == nb:
                    DFLX_im1 = wp(0.)
                    CFLX = wp(0.)
                    EFLX = wp(0.)
                if j == ny + nb - 1:
                    DFLX_jp1 = wp(0.)
                    CFLX_jp1 = wp(0.)
                    EFLX_im1_jp1 = wp(0.)

                dUFLXdt[i, j, k], dUFLXdt_TURB[i, j, k] = add_up_tendencies(
                    # 3D
                    UFLX[i, j, k],
                    UFLX[i - 1, j, k],
                    UFLX[i + 1, j, k],
                    UFLX[i, j - 1, k],
                    UFLX[i, j + 1, k],
                    VWIND[i, j, k],
                    VWIND[i - 1, j, k],
                    VWIND[i, j + 1, k],
                    VWIND[i - 1, j + 1, k],
                    UWIND[i, j, k],
                    UWIND[i - 1, j, k],
                    UWIND[i + 1, j, k],
                    UWIND[i, j - 1, k],
                    UWIND[i, j + 1, k],
                    UWIND[i - 1, j - 1, k],
                    UWIND[i - 1, j + 1, k],
                    UWIND[i + 1, j - 1, k],
                    UWIND[i + 1, j + 1, k],
                    BFLX,
                    BFLX_im1,
                    CFLX,
                    CFLX_jp1,
                    DFLX_im1,
                    DFLX_jp1,
                    EFLX,
                    EFLX_im1_jp1,
                    PHI[i, j, k],
                    PHI[i - 1, j, k],
                    POTT[i, j, k],
                    POTT[i - 1, j, k],
                    PVTF[i, j, k],
                    PVTF[i - 1, j, k],
                    PVTFVB[i, j, k],
                    PVTFVB[i - 1, j, k],
                    PVTFVB[i - 1, j, k + 1],
                    PVTFVB[i, j, k + 1],
                    WWIND_UWIND[i, j, k],
                    WWIND_UWIND[i, j, k + 1],
                    PHIVB[i, j, k],
                    PHIVB[i - 1, j, k],
                    PHIVB[i, j - 1, k],
                    PHIVB[i, j + 1, k],
                    PHIVB[i - 1, j - 1, k],
                    PHIVB[i - 1, j + 1, k],
                    PHIVB[i, j, k + 1],
                    PHIVB[i - 1, j, k + 1],
                    PHIVB[i, j - 1, k + 1],
                    PHIVB[i, j + 1, k + 1],
                    PHIVB[i - 1, j - 1, k + 1],
                    PHIVB[i - 1, j + 1, k + 1],
                    RHO[i, j, k],
                    RHO[i - 1, j, k],
                    RHO[i, j - 1, k],
                    RHO[i, j + 1, k],
                    RHO[i - 1, j - 1, k],
                    RHO[i - 1, j + 1, k],
                    SMOMXFLX[i, j, 0],
                    SMOMXFLX[i - 1, j, 0],
                    SMOMXFLX[i, j - 1, 0],
                    SMOMXFLX[i, j + 1, 0],
                    SMOMXFLX[i - 1, j - 1, 0],
                    SMOMXFLX[i - 1, j + 1, 0],
                    KMOM_dUWINDdz[i, j, k],
                    KMOM_dUWINDdz[i, j, k + 1],

                    # 2D
                    COLP[i, j, 0],
                    COLP[i - 1, j, 0],
                    # GR horizontal
                    corf_is[i, j, 0],
                    lat_is_rad[i, j, 0],
                    dlon_rad[i, j, 0],
                    dlat_rad[i, j, 0],
                    dyis[i, j, 0],
                    # GR vertical
                    dsigma[0, 0, k],
                    sigma_vb[0, 0, k],
                    sigma_vb[0, 0, k + 1],
                    UVFLX_dif_coef[0, 0, k],
                    k,
                    j)
Ejemplo n.º 13
0
def launch_numba_cpu_prep_adv(WWIND_UWIND, WWIND_VWIND,
                                UWIND, VWIND, WWIND,
                                UFLX_3D, VFLX_3D,
                                CFLX, QFLX, DFLX, EFLX,
                                SFLX, TFLX, BFLX, RFLX,
                                COLP, COLP_NEW,
                                KMOM_dUWINDdz, KMOM_dVWINDdz,
                                KMOM, PHI, RHOVB,
                                A, dsigma):

    ###########################################################################
    if i_UVFLX_vert_adv:

        for i in prange(nb,nxs+nb):
            for j in range(nb,ny+nb):
                for k in range(wp_int(0),nzs):
                    WWIND_UWIND[i  ,j  ,k  ] = \
                        interp_WWIND_UVWIND(
                UWIND     [i  ,j  ,k  ], UWIND     [i  ,j  ,k-1],
                WWIND     [i  ,j  ,k  ], WWIND     [i-1,j  ,k  ],
                WWIND     [i  ,j-1,k  ], WWIND     [i  ,j+1,k  ],
                WWIND     [i-1,j-1,k  ], WWIND     [i-1,j+1,k  ], 
                COLP_NEW  [i  ,j  ,0  ], COLP_NEW  [i-1,j  ,0  ],
                COLP_NEW  [i  ,j-1,0  ], COLP_NEW  [i  ,j+1,0  ],
                COLP_NEW  [i-1,j-1,0  ], COLP_NEW  [i-1,j+1,0  ], 
                A         [i  ,j  ,0  ], A         [i-1,j  ,0  ],
                A         [i  ,j-1,0  ], A         [i  ,j+1,0  ],
                A         [i-1,j-1,0  ], A         [i-1,j+1,0  ], 
                dsigma    [0  ,0  ,k  ], dsigma    [0  ,0  ,k-1],
                True, j, ny, k)

        for i in prange(nb,nx+nb):
            for j in range(nb,nys+nb):
                for k in range(wp_int(0),nzs):
                    WWIND_VWIND[i  ,j  ,k  ] = \
                        interp_WWIND_UVWIND(
                VWIND     [i  ,j  ,k  ], VWIND     [i  ,j  ,k-1],
                WWIND     [i  ,j  ,k  ], WWIND     [i  ,j-1,k  ],
                WWIND     [i-1,j  ,k  ], WWIND     [i+1,j  ,k  ],
                WWIND     [i-1,j-1,k  ], WWIND     [i+1,j-1,k  ], 
                COLP_NEW  [i  ,j  ,0  ], COLP_NEW  [i  ,j-1,0  ],
                COLP_NEW  [i-1,j  ,0  ], COLP_NEW  [i+1,j  ,0  ],
                COLP_NEW  [i-1,j-1,0  ], COLP_NEW  [i+1,j-1,0  ], 
                A         [i  ,j  ,0  ], A         [i  ,j-1,0  ],
                A         [i-1,j  ,0  ], A         [i+1,j  ,0  ],
                A         [i-1,j-1,0  ], A         [i+1,j-1,0  ], 
                dsigma    [0  ,0  ,k  ], dsigma    [0  ,0  ,k-1],
                False, i, nx, k)

    ###########################################################################
    if i_UVFLX_vert_turb:
        for i in prange(nb,nxs+nb):
            for j in range(nb,ny+nb):
                for k in range(wp_int(0),nzs):
                    KMOM_dUWINDdz[i  ,j  ,k  ] = \
                interp_KMOM_dUVWINDdz(
                UWIND     [i  ,j  ,k  ], UWIND     [i  ,j  ,k-1],
                KMOM      [i  ,j  ,k  ], KMOM      [i-1,j  ,k  ],
                KMOM      [i  ,j-1,k  ], KMOM      [i  ,j+1,k  ],
                KMOM      [i-1,j-1,k  ], KMOM      [i-1,j+1,k  ], 
                RHOVB     [i  ,j  ,k  ], RHOVB     [i-1,j  ,k  ],
                RHOVB     [i  ,j-1,k  ], RHOVB     [i  ,j+1,k  ],
                RHOVB     [i-1,j-1,k  ], RHOVB     [i-1,j+1,k  ], 
                PHI       [i  ,j  ,k  ], PHI       [i-1,j  ,k  ],
                PHI       [i  ,j-1,k  ], PHI       [i  ,j+1,k  ],
                PHI       [i-1,j-1,k  ], PHI       [i-1,j+1,k  ], 
                PHI       [i  ,j  ,k-1], PHI       [i-1,j  ,k-1],
                PHI       [i  ,j-1,k-1], PHI       [i  ,j+1,k-1],
                PHI       [i-1,j-1,k-1], PHI       [i-1,j+1,k-1], 
                COLP      [i  ,j  ,0  ], COLP      [i-1,j  ,0  ],
                COLP      [i  ,j-1,0  ], COLP      [i  ,j+1,0  ],
                COLP      [i-1,j-1,0  ], COLP      [i-1,j+1,0  ], 
                A         [i  ,j  ,0  ], A         [i-1,j  ,0  ],
                A         [i  ,j-1,0  ], A         [i  ,j+1,0  ],
                A         [i-1,j-1,0  ], A         [i-1,j+1,0  ], 
                dsigma    [0  ,0  ,k  ], dsigma    [0  ,0  ,k-1],
                True, j, ny, k)

        for i in prange(nb,nx+nb):
            for j in range(nb,nys+nb):
                for k in range(wp_int(0),nzs):
                    KMOM_dVWINDdz[i  ,j  ,k  ] = \
                interp_KMOM_dUVWINDdz(
                VWIND     [i  ,j  ,k  ], VWIND     [i  ,j  ,k-1],
                KMOM      [i  ,j  ,k  ], KMOM      [i  ,j-1,k  ],
                KMOM      [i-1,j  ,k  ], KMOM      [i+1,j  ,k  ],
                KMOM      [i-1,j-1,k  ], KMOM      [i+1,j-1,k  ], 
                RHOVB     [i  ,j  ,k  ], RHOVB     [i  ,j-1,k  ],
                RHOVB     [i-1,j  ,k  ], RHOVB     [i+1,j  ,k  ],
                RHOVB     [i-1,j-1,k  ], RHOVB     [i+1,j-1,k  ], 
                PHI       [i  ,j  ,k  ], PHI       [i  ,j-1,k  ],
                PHI       [i-1,j  ,k  ], PHI       [i+1,j  ,k  ],
                PHI       [i-1,j-1,k  ], PHI       [i+1,j-1,k  ], 
                PHI       [i  ,j  ,k-1], PHI       [i  ,j-1,k-1],
                PHI       [i-1,j  ,k-1], PHI       [i+1,j  ,k-1],
                PHI       [i-1,j-1,k-1], PHI       [i+1,j-1,k-1], 
                COLP      [i  ,j  ,0  ], COLP      [i  ,j-1,0  ],
                COLP      [i-1,j  ,0  ], COLP      [i+1,j  ,0  ],
                COLP      [i-1,j-1,0  ], COLP      [i+1,j-1,0  ], 
                A         [i  ,j  ,0  ], A         [i  ,j-1,0  ],
                A         [i-1,j  ,0  ], A         [i+1,j  ,0  ],
                A         [i-1,j-1,0  ], A         [i+1,j-1,0  ], 
                dsigma    [0  ,0  ,k  ], dsigma    [0  ,0  ,k-1],
                False, i, nx, k)


    ###########################################################################
    if i_UVFLX_hor_adv:

        for i in prange(nb,nxs+nb):
            for j in range(nb,nys+nb):
                for k in range(wp_int(0),nz):
                    UFLX            = UFLX_3D[i  ,j  ,k  ]
                    UFLX_im1        = UFLX_3D[i-1,j  ,k  ]
                    UFLX_im1_jm1    = UFLX_3D[i-1,j-1,k  ]
                    UFLX_im1_jp1    = UFLX_3D[i-1,j+1,k  ]
                    UFLX_ip1        = UFLX_3D[i+1,j  ,k  ]
                    UFLX_ip1_jm1    = UFLX_3D[i+1,j-1,k  ]
                    UFLX_ip1_jp1    = UFLX_3D[i+1,j+1,k  ]
                    UFLX_jm1        = UFLX_3D[i  ,j-1,k  ]
                    UFLX_jp1        = UFLX_3D[i  ,j+1,k  ]

                    VFLX            = VFLX_3D[i  ,j  ,k  ]
                    VFLX_im1        = VFLX_3D[i-1,j  ,k  ]
                    VFLX_im1_jm1    = VFLX_3D[i-1,j-1,k  ]
                    VFLX_im1_jp1    = VFLX_3D[i-1,j+1,k  ]
                    VFLX_ip1        = VFLX_3D[i+1,j  ,k  ]
                    VFLX_ip1_jm1    = VFLX_3D[i+1,j-1,k  ]
                    VFLX_ip1_jp1    = VFLX_3D[i+1,j+1,k  ]
                    VFLX_jm1        = VFLX_3D[i  ,j-1,k  ]
                    VFLX_jp1        = VFLX_3D[i  ,j+1,k  ]

                    if i < nxs+nb and j < nys+nb:
                        CFLX[i,j,k],QFLX[i,j,k] = \
                            calc_momentum_fluxes_isjs(
                                UFLX, UFLX_im1,
                                UFLX_im1_jm1, UFLX_im1_jp1,
                                UFLX_ip1, UFLX_ip1_jm1,
                                UFLX_ip1_jp1, UFLX_jm1,
                                UFLX_jp1,
                                VFLX, VFLX_im1,
                                VFLX_im1_jm1, VFLX_im1_jp1,
                                VFLX_ip1, VFLX_ip1_jm1,
                                VFLX_ip1_jp1, VFLX_jm1,
                                VFLX_jp1)

                    if i < nx+nb and j < nys+nb:
                        DFLX[i,j,k],EFLX[i,j,k] = \
                            calc_momentum_fluxes_ijs(
                                UFLX, UFLX_im1,
                                UFLX_im1_jm1, UFLX_im1_jp1,
                                UFLX_ip1, UFLX_ip1_jm1,
                                UFLX_ip1_jp1, UFLX_jm1,
                                UFLX_jp1,
                                VFLX, VFLX_im1,
                                VFLX_im1_jm1, VFLX_im1_jp1,
                                VFLX_ip1, VFLX_ip1_jm1,
                                VFLX_ip1_jp1, VFLX_jm1,
                                VFLX_jp1)

                    if i < nxs+nb and j < ny+nb:
                        SFLX[i,j,k],TFLX[i,j,k] = \
                            calc_momentum_fluxes_isj(
                                UFLX, UFLX_im1,
                                UFLX_im1_jm1, UFLX_im1_jp1,
                                UFLX_ip1, UFLX_ip1_jm1,
                                UFLX_ip1_jp1, UFLX_jm1,
                                UFLX_jp1,
                                VFLX, VFLX_im1,
                                VFLX_im1_jm1, VFLX_im1_jp1,
                                VFLX_ip1, VFLX_ip1_jm1,
                                VFLX_ip1_jp1, VFLX_jm1,
                                VFLX_jp1)


                    if i < nx+nb and j < ny+nb:
                        BFLX[i,j,k],RFLX[i,j,k] = \
                            calc_momentum_fluxes_ij(
                                UFLX, UFLX_im1,
                                UFLX_im1_jm1, UFLX_im1_jp1,
                                UFLX_ip1, UFLX_ip1_jm1,
                                UFLX_ip1_jp1, UFLX_jm1,
                                UFLX_jp1,
                                VFLX, VFLX_im1,
                                VFLX_im1_jm1, VFLX_im1_jp1,
                                VFLX_ip1, VFLX_ip1_jm1,
                                VFLX_ip1_jp1, VFLX_jm1,
                                VFLX_jp1)
Ejemplo n.º 14
0
def add_up_tendencies_py(
        VFLX, VFLX_im1, VFLX_ip1, VFLX_jm1, VFLX_jp1, UWIND, UWIND_jm1,
        UWIND_ip1, UWIND_ip1_jm1, VWIND, VWIND_jm1, VWIND_jp1, VWIND_im1,
        VWIND_ip1, VWIND_jm1_im1, VWIND_jm1_ip1, VWIND_jp1_im1, VWIND_jp1_ip1,
        RFLX, RFLX_jm1, QFLX, QFLX_ip1, SFLX_jm1, SFLX_ip1, TFLX, TFLX_jm1_ip1,
        PHI, PHI_jm1, POTT, POTT_jm1, PVTF, PVTF_jm1, PVTFVB, PVTFVB_jm1,
        PVTFVB_jm1_kp1, PVTFVB_kp1, WWIND_VWIND, WWIND_VWIND_kp1, PHIVB,
        PHIVB_jm1, PHIVB_im1, PHIVB_ip1, PHIVB_im1_jm1, PHIVB_ip1_jm1,
        PHIVB_kp1, PHIVB_kp1_jm1, PHIVB_kp1_im1, PHIVB_kp1_ip1,
        PHIVB_kp1_im1_jm1, PHIVB_kp1_ip1_jm1, RHO, RHO_jm1, RHO_im1, RHO_ip1,
        RHO_im1_jm1, RHO_ip1_jm1, SMOMYFLX, SMOMYFLX_jm1, SMOMYFLX_im1,
        SMOMYFLX_ip1, SMOMYFLX_im1_jm1, SMOMYFLX_ip1_jm1, KMOM_dVWINDdz,
        KMOM_dVWINDdz_kp1, COLP, COLP_jm1, corf, corf_jm1, lat_rad,
        lat_rad_jm1, dlon_rad, dlat_rad, dxjs, dsigma, sigma_vb, sigma_vb_kp1,
        UVFLX_dif_coef, k, i):
    """
    Compute and add up all tendency contributions of VFLUX.
    """

    dVFLXdt = wp(0.)

    if i_UVFLX_main_switch:
        # HORIZONTAL ADVECTION
        if i_UVFLX_hor_adv:
            #pass
            dVFLXdt = dVFLXdt + UVFLX_hor_adv(
                VWIND, VWIND_jm1, VWIND_jp1, VWIND_im1, VWIND_ip1,
                VWIND_jm1_im1, VWIND_jm1_ip1, VWIND_jp1_im1, VWIND_jp1_ip1,
                RFLX, RFLX_jm1, QFLX, QFLX_ip1, SFLX_jm1, SFLX_ip1, TFLX,
                TFLX_jm1_ip1, wp(-1.))
        # VERTICAL ADVECTION
        if i_UVFLX_vert_adv:
            dVFLXdt = dVFLXdt + ((WWIND_VWIND - WWIND_VWIND_kp1) / dsigma)
        # VERTICAL TURBULENT TRANSPORT
        if i_UVFLX_vert_turb:

            ALTVB_js = interp_VAR_ds(PHIVB, PHIVB_jm1, PHIVB_im1, PHIVB_ip1,
                                     PHIVB_im1_jm1, PHIVB_ip1_jm1, False, i,
                                     nx) / con_g
            ALTVB_kp1_js = interp_VAR_ds(
                PHIVB_kp1, PHIVB_kp1_jm1, PHIVB_kp1_im1, PHIVB_kp1_ip1,
                PHIVB_kp1_im1_jm1, PHIVB_kp1_ip1_jm1, False, i, nx) / con_g
            RHO_js = interp_VAR_ds(RHO, RHO_jm1, RHO_im1, RHO_ip1, RHO_im1_jm1,
                                   RHO_ip1_jm1, False, i, nx)
            SMOMYFLX_js = interp_VAR_ds(SMOMYFLX, SMOMYFLX_jm1, SMOMYFLX_im1,
                                        SMOMYFLX_ip1, SMOMYFLX_im1_jm1,
                                        SMOMYFLX_ip1_jm1, False, i, nx)

            if k == wp_int(0):
                dVFLXdt_TURB = ((wp(0.) - KMOM_dVWINDdz_kp1) /
                                ((ALTVB_js - ALTVB_kp1_js) * RHO_js))
            elif k == wp_int(nz - 1):
                dVFLXdt_TURB = ((KMOM_dVWINDdz + SMOMYFLX_js) /
                                ((ALTVB_js - ALTVB_kp1_js) * RHO_js))
                #dVFLXdt_TURB = wp(0.)
            else:
                dVFLXdt_TURB = ((KMOM_dVWINDdz - KMOM_dVWINDdz_kp1) /
                                ((ALTVB_js - ALTVB_kp1_js) * RHO_js))

            #dVFLXdt_TURB = wp(0.)
            dVFLXdt = dVFLXdt + dVFLXdt_TURB

        # CORIOLIS AND SPHERICAL GRID CONVERSION
        if i_UVFLX_coriolis:
            dVFLXdt = dVFLXdt + coriolis_and_spherical_VWIND(
                COLP, COLP_jm1, UWIND, UWIND_jm1, UWIND_ip1, UWIND_ip1_jm1,
                corf, corf_jm1, lat_rad, lat_rad_jm1, dlon_rad, dlat_rad)

        # PRESSURE GRADIENT
        if i_UVFLX_pre_grad:
            dVFLXdt = dVFLXdt + pre_grad(
                PHI, PHI_jm1, COLP, COLP_jm1, POTT, POTT_jm1, PVTF, PVTF_jm1,
                PVTFVB, PVTFVB_jm1, PVTFVB_jm1_kp1, PVTFVB_kp1, dsigma,
                sigma_vb, sigma_vb_kp1, dxjs)
        # NUMERICAL HORIZONTAL DIFUSION
        if i_UVFLX_num_dif and (UVFLX_dif_coef > wp(0.)):
            dVFLXdt = dVFLXdt + num_dif(VFLX, VFLX_im1, VFLX_ip1, VFLX_jm1,
                                        VFLX_jp1, UVFLX_dif_coef)

    return (dVFLXdt, dVFLXdt_TURB)
Ejemplo n.º 15
0
def launch_numba_cpu_main(dVFLXdt, VFLX, UWIND, VWIND, RFLX_3D, SFLX_3D,
                          TFLX_3D, QFLX_3D, PHI, PHIVB, COLP, POTT, PVTF,
                          PVTFVB, WWIND_VWIND, KMOM_dVWINDdz, RHO,
                          dVFLXdt_TURB, SMOMYFLX, corf, lat_rad, dlon_rad,
                          dlat_rad, dxjs, dsigma, sigma_vb, UVFLX_dif_coef):

    for i in prange(nb, nx + nb):
        for j in range(nb, nys + nb):
            for k in range(wp_int(0), nz):
                # Prepare momentum fluxes and set boundary conditions if
                # necessary
                RFLX = RFLX_3D[i, j, k]
                QFLX = QFLX_3D[i, j, k]
                TFLX = TFLX_3D[i, j, k]
                RFLX_jm1 = RFLX_3D[i, j - 1, k]
                SFLX_jm1 = SFLX_3D[i, j - 1, k]

                # BCx is
                if i == nxs - 1:
                    QFLX_ip1 = QFLX_3D[1, j, k]
                    TFLX_ip1_jm1 = TFLX_3D[1, j - 1, k]
                    SFLX_ip1 = SFLX_3D[1, j, k]
                else:
                    QFLX_ip1 = QFLX_3D[i + 1, j, k]
                    TFLX_ip1_jm1 = TFLX_3D[i + 1, j - 1, k]
                    SFLX_ip1 = SFLX_3D[i + 1, j, k]

                #dVFLXdt[i  ,j  ,k] = add_up_tendencies(
                dVFLXdt[i, j, k], dVFLXdt_TURB[i, j, k] = add_up_tendencies(
                    # 3D
                    VFLX[i, j, k],
                    VFLX[i - 1, j, k],
                    VFLX[i + 1, j, k],
                    VFLX[i, j - 1, k],
                    VFLX[i, j + 1, k],
                    UWIND[i, j, k],
                    UWIND[i, j - 1, k],
                    UWIND[i + 1, j, k],
                    UWIND[i + 1, j - 1, k],
                    VWIND[i, j, k],
                    VWIND[i, j - 1, k],
                    VWIND[i, j + 1, k],
                    VWIND[i - 1, j, k],
                    VWIND[i + 1, j, k],
                    VWIND[i - 1, j - 1, k],
                    VWIND[i + 1, j - 1, k],
                    VWIND[i - 1, j + 1, k],
                    VWIND[i + 1, j + 1, k],
                    RFLX,
                    RFLX_jm1,
                    QFLX,
                    QFLX_ip1,
                    SFLX_jm1,
                    SFLX_ip1,
                    TFLX,
                    TFLX_ip1_jm1,
                    PHI[i, j, k],
                    PHI[i, j - 1, k],
                    POTT[i, j, k],
                    POTT[i, j - 1, k],
                    PVTF[i, j, k],
                    PVTF[i, j - 1, k],
                    PVTFVB[i, j, k],
                    PVTFVB[i, j - 1, k],
                    PVTFVB[i, j - 1, k + 1],
                    PVTFVB[i, j, k + 1],
                    WWIND_VWIND[i, j, k],
                    WWIND_VWIND[i, j, k + 1],
                    PHIVB[i, j, k],
                    PHIVB[i, j - 1, k],
                    PHIVB[i - 1, j, k],
                    PHIVB[i + 1, j, k],
                    PHIVB[i - 1, j - 1, k],
                    PHIVB[i + 1, j - 1, k],
                    PHIVB[i, j, k + 1],
                    PHIVB[i, j - 1, k + 1],
                    PHIVB[i - 1, j, k + 1],
                    PHIVB[i + 1, j, k + 1],
                    PHIVB[i - 1, j - 1, k + 1],
                    PHIVB[i + 1, j - 1, k + 1],
                    RHO[i, j, k],
                    RHO[i, j - 1, k],
                    RHO[i - 1, j, k],
                    RHO[i + 1, j, k],
                    RHO[i - 1, j - 1, k],
                    RHO[i + 1, j - 1, k],
                    SMOMYFLX[i, j, 0],
                    SMOMYFLX[i, j - 1, 0],
                    SMOMYFLX[i - 1, j, 0],
                    SMOMYFLX[i + 1, j, 0],
                    SMOMYFLX[i - 1, j - 1, 0],
                    SMOMYFLX[i + 1, j - 1, 0],
                    KMOM_dVWINDdz[i, j, k],
                    KMOM_dVWINDdz[i, j, k + 1],

                    # 2D
                    COLP[i, j, 0],
                    COLP[i, j - 1, 0],
                    # GR horizontal
                    corf[i, j, 0],
                    corf[i, j - 1, 0],
                    lat_rad[i, j, 0],
                    lat_rad[i, j - 1, 0],
                    dlon_rad[i, j, 0],
                    dlat_rad[i, j, 0],
                    dxjs[i, j, 0],
                    # GR vertical
                    dsigma[0, 0, k],
                    sigma_vb[0, 0, k],
                    sigma_vb[0, 0, k + 1],
                    UVFLX_dif_coef[0, 0, k],
                    k,
                    i)
Ejemplo n.º 16
0
                            moist_dif_coef)
from io_constants import con_rE, con_omega
from io_restart import load_restart_grid
from io_initial_conditions import set_up_sigma_levels
from misc_utilities import Timer
###############################################################################

###############################################################################
# GLOBAL SETTINGS
###############################################################################
# vertical extent of cuda shared arrays
# (must be called before setting nz = wp_int(nz))
shared_nz = nz 

# domain extent to be imported in modules
nz = wp_int(nz)
nzs = wp_int(nz+1)
nx = wp_int((lon1_deg - lon0_deg)/dlon_deg)
nxs = wp_int(nx+1)
ny = wp_int((lat1_deg - lat0_deg)/dlat_deg)
nys = wp_int(ny+1)
nb = wp_int(nb)

###############################################################################
# GPU COMPUTATION SETTINGS
###############################################################################
if gpu_enable:
    if nz > 32:
        raise NotImplementedError('More than 32 vertical levels '+
                'not yet implemented for GPU.')
Ejemplo n.º 17
0
def make_timestep_cpu(COLP, COLP_OLD, UWIND, UWIND_OLD, dUFLXdt, VWIND,
                      VWIND_OLD, dVFLXdt, POTT, POTT_OLD, dPOTTdt, QV, QV_OLD,
                      dQVdt, QC, QC_OLD, dQCdt, A, dt):

    for i in prange(nb, nxs + nb):
        for j in range(nb, nys + nb):

            COLP_ = COLP[i, j, 0]
            COLP_im1 = COLP[i - 1, j, 0]
            COLP_ip1 = COLP[i + 1, j, 0]
            COLP_jm1 = COLP[i, j - 1, 0]
            COLP_jp1 = COLP[i, j + 1, 0]
            COLP_im1_jm1 = COLP[i - 1, j - 1, 0]
            COLP_im1_jp1 = COLP[i - 1, j + 1, 0]
            COLP_ip1_jm1 = COLP[i + 1, j - 1, 0]
            COLP_ip1_jp1 = COLP[i + 1, j + 1, 0]

            COLP_OLD_ = COLP_OLD[i, j, 0]
            COLP_OLD_im1 = COLP_OLD[i - 1, j, 0]
            COLP_OLD_ip1 = COLP_OLD[i + 1, j, 0]
            COLP_OLD_jm1 = COLP_OLD[i, j - 1, 0]
            COLP_OLD_jp1 = COLP_OLD[i, j + 1, 0]
            COLP_OLD_im1_jm1 = COLP_OLD[i - 1, j - 1, 0]
            COLP_OLD_im1_jp1 = COLP_OLD[i - 1, j + 1, 0]
            COLP_OLD_ip1_jm1 = COLP_OLD[i + 1, j - 1, 0]
            COLP_OLD_ip1_jp1 = COLP_OLD[i + 1, j + 1, 0]

            A_ = A[i, j, 0]
            A_im1 = A[i - 1, j, 0]
            A_ip1 = A[i + 1, j, 0]
            A_jm1 = A[i, j - 1, 0]
            A_jp1 = A[i, j + 1, 0]
            A_im1_jm1 = A[i - 1, j - 1, 0]
            A_im1_jp1 = A[i - 1, j + 1, 0]
            A_ip1_jm1 = A[i + 1, j - 1, 0]
            A_ip1_jp1 = A[i + 1, j + 1, 0]

            ## UWIND
            COLPA_is = interp_COLPA_is(COLP_, COLP_im1, COLP_jm1, COLP_jp1,
                                       COLP_im1_jp1, COLP_im1_jm1, A_, A_im1,
                                       A_jm1, A_jp1, A_im1_jp1, A_im1_jm1, j)
            COLPA_OLD_is = interp_COLPA_is(COLP_OLD_, COLP_OLD_im1,
                                           COLP_OLD_jm1, COLP_OLD_jp1,
                                           COLP_OLD_im1_jp1, COLP_OLD_im1_jm1,
                                           A_, A_im1, A_jm1, A_jp1, A_im1_jp1,
                                           A_im1_jm1, j)

            # VWIND
            COLPA_js = interp_COLPA_js(COLP_, COLP_jm1, COLP_im1, COLP_ip1,
                                       COLP_ip1_jm1, COLP_im1_jm1, A_, A_jm1,
                                       A_im1, A_ip1, A_ip1_jm1, A_im1_jm1)
            COLPA_OLD_js = interp_COLPA_js(COLP_OLD_, COLP_OLD_jm1,
                                           COLP_OLD_im1, COLP_OLD_ip1,
                                           COLP_OLD_ip1_jm1, COLP_OLD_im1_jm1,
                                           A_, A_jm1, A_im1, A_ip1, A_ip1_jm1,
                                           A_im1_jm1)

            for k in range(wp_int(0), nz):
                UWIND_OLD_ = UWIND_OLD[i, j, k]
                dUFLXdt_ = dUFLXdt[i, j, k]
                VWIND_OLD_ = VWIND_OLD[i, j, k]
                dVFLXdt_ = dVFLXdt[i, j, k]
                POTT_OLD_ = POTT_OLD[i, j, k]
                dPOTTdt_ = dPOTTdt[i, j, k]
                QV_OLD_ = QV_OLD[i, j, k]
                dQVdt_ = dQVdt[i, j, k]
                QC_OLD_ = QC_OLD[i, j, k]
                dQCdt_ = dQCdt[i, j, k]

                UWIND[i, j, k] = euler_forward_pw(UWIND_OLD_, dUFLXdt_,
                                                  COLPA_is, COLPA_OLD_is, dt)
                VWIND[i, j, k] = euler_forward_pw(VWIND_OLD_, dVFLXdt_,
                                                  COLPA_js, COLPA_OLD_js, dt)
                POTT[i, j, k] = euler_forward_pw(POTT_OLD_, dPOTTdt_, COLP_,
                                                 COLP_OLD_, dt)
                QV[i, j, k] = euler_forward_pw(QV_OLD_, dQVdt_, COLP_,
                                               COLP_OLD_, dt)
                QC[i, j, k] = euler_forward_pw(QC_OLD_, dQCdt_, COLP_,
                                               COLP_OLD_, dt)
Ejemplo n.º 18
0
def add_up_tendencies_py(
        UFLX, UFLX_im1, UFLX_ip1, UFLX_jm1, UFLX_jp1, VWIND, VWIND_im1,
        VWIND_jp1, VWIND_im1_jp1, UWIND, UWIND_im1, UWIND_ip1, UWIND_jm1,
        UWIND_jp1, UWIND_im1_jm1, UWIND_im1_jp1, UWIND_ip1_jm1, UWIND_ip1_jp1,
        BFLX, BFLX_im1, CFLX, CFLX_jp1, DFLX_im1, DFLX_jp1, EFLX, EFLX_im1_jp1,
        PHI, PHI_im1, POTT, POTT_im1, PVTF, PVTF_im1, PVTFVB, PVTFVB_im1,
        PVTFVB_im1_kp1, PVTFVB_kp1, WWIND_UWIND, WWIND_UWIND_kp1, PHIVB,
        PHIVB_im1, PHIVB_jm1, PHIVB_jp1, PHIVB_jm1_im1, PHIVB_jp1_im1,
        PHIVB_kp1, PHIVB_kp1_im1, PHIVB_kp1_jm1, PHIVB_kp1_jp1,
        PHIVB_kp1_jm1_im1, PHIVB_kp1_jp1_im1, RHO, RHO_im1, RHO_jm1, RHO_jp1,
        RHO_jm1_im1, RHO_jp1_im1, SMOMXFLX, SMOMXFLX_im1, SMOMXFLX_jm1,
        SMOMXFLX_jp1, SMOMXFLX_jm1_im1, SMOMXFLX_jp1_im1, KMOM_dUWINDdz,
        KMOM_dUWINDdz_kp1, COLP, COLP_im1, corf_is, lat_is_rad, dlon_rad,
        dlat_rad, dyis, dsigma, sigma_vb, sigma_vb_kp1, UVFLX_dif_coef, k, j):
    """
    Compute and add up all tendency contributions of UFLUX.
    """

    dUFLXdt = wp(0.)

    if i_UVFLX_main_switch:
        # HORIZONTAL ADVECTION
        if i_UVFLX_hor_adv:
            dUFLXdt = dUFLXdt + UVFLX_hor_adv(
                UWIND, UWIND_im1, UWIND_ip1, UWIND_jm1, UWIND_jp1,
                UWIND_im1_jm1, UWIND_im1_jp1, UWIND_ip1_jm1, UWIND_ip1_jp1,
                BFLX, BFLX_im1, CFLX, CFLX_jp1, DFLX_im1, DFLX_jp1, EFLX,
                EFLX_im1_jp1, wp(1.))
        # VERTICAL ADVECTION
        if i_UVFLX_vert_adv:
            dUFLXdt = dUFLXdt + ((WWIND_UWIND - WWIND_UWIND_kp1) / dsigma)
        # VERTICAL TURBULENT TRANSPORT
        if i_UVFLX_vert_turb:
            ALTVB_is = interp_VAR_ds(PHIVB, PHIVB_im1, PHIVB_jm1, PHIVB_jp1,
                                     PHIVB_jm1_im1, PHIVB_jp1_im1, True, j,
                                     ny) / con_g
            ALTVB_kp1_is = interp_VAR_ds(
                PHIVB_kp1, PHIVB_kp1_im1, PHIVB_kp1_jm1, PHIVB_kp1_jp1,
                PHIVB_kp1_jm1_im1, PHIVB_kp1_jp1_im1, True, j, ny) / con_g
            RHO_is = interp_VAR_ds(RHO, RHO_im1, RHO_jm1, RHO_jp1, RHO_jm1_im1,
                                   RHO_jp1_im1, True, j, ny)
            SMOMXFLX_is = interp_VAR_ds(SMOMXFLX, SMOMXFLX_im1, SMOMXFLX_jm1,
                                        SMOMXFLX_jp1, SMOMXFLX_jm1_im1,
                                        SMOMXFLX_jp1_im1, True, j, ny)

            if k == wp_int(0):
                dUFLXdt_TURB = ((wp(0.) - KMOM_dUWINDdz_kp1) /
                                ((ALTVB_is - ALTVB_kp1_is) * RHO_is))
            elif k == wp_int(nz - 1):
                dUFLXdt_TURB = ((KMOM_dUWINDdz + SMOMXFLX_is) /
                                ((ALTVB_is - ALTVB_kp1_is) * RHO_is))
                #dUFLXdt_TURB = wp(0.)
            else:
                dUFLXdt_TURB = ((KMOM_dUWINDdz - KMOM_dUWINDdz_kp1) /
                                ((ALTVB_is - ALTVB_kp1_is) * RHO_is))

            #dUFLXdt_TURB = wp(0.)
            dUFLXdt = dUFLXdt + dUFLXdt_TURB

        # CORIOLIS AND SPHERICAL GRID CONVERSION
        if i_UVFLX_coriolis:
            dUFLXdt = dUFLXdt + coriolis_and_spherical_UWIND(
                COLP, COLP_im1, VWIND, VWIND_im1, VWIND_jp1, VWIND_im1_jp1,
                UWIND, UWIND_im1, UWIND_ip1, corf_is, lat_is_rad, dlon_rad,
                dlat_rad)
        # PRESSURE GRADIENT
        if i_UVFLX_pre_grad:
            dUFLXdt = dUFLXdt + pre_grad(
                PHI, PHI_im1, COLP, COLP_im1, POTT, POTT_im1, PVTF, PVTF_im1,
                PVTFVB, PVTFVB_im1, PVTFVB_im1_kp1, PVTFVB_kp1, dsigma,
                sigma_vb, sigma_vb_kp1, dyis)
        # NUMERICAL HORIZONTAL DIFUSION
        if i_UVFLX_num_dif and (UVFLX_dif_coef > wp(0.)):
            dUFLXdt = dUFLXdt + num_dif(UFLX, UFLX_im1, UFLX_ip1, UFLX_jm1,
                                        UFLX_jp1, UVFLX_dif_coef)

    return (dUFLXdt, dUFLXdt_TURB)