diff --git a/setup.py b/setup.py index 0563d40..f80706e 100644 --- a/setup.py +++ b/setup.py @@ -83,6 +83,19 @@ extra_f90_compile_args=f90flags, f2py_options=['--quiet']) +ext_interpI_pl = Extension(name = 'Interp_pressure_lev', + sources = ['xcape/Interp_pressure_lev.pyf', + 'xcape/Interp_pressure_lev.f90'], + extra_f90_compile_args=f90flags, + f2py_options=['--quiet']) + +ext_interpI_ml = Extension(name = 'Interp_model_lev', + sources = ['xcape/Interp_model_lev.pyf', + 'xcape/Interp_model_lev.f90'], + extra_f90_compile_args=f90flags, + f2py_options=['--quiet']) + + DESCRIPTION = "Fast convective parameters for numpy, dask, and xarray" def readme(): with open('README.rst') as f: @@ -105,6 +118,7 @@ def readme(): ext_modules = [ext_cape_ml, ext_cape_pl, ext_bunkers_ml, ext_bunkers_pl, ext_srh_ml, ext_srh_pl, - ext_stdh_ml, ext_stdh_pl] + ext_stdh_ml, ext_stdh_pl, + ext_interpI_pl, ext_interpI_ml] ) diff --git a/xcape/Indices_calc.py b/xcape/Indices_calc.py new file mode 100644 index 0000000..635b41e --- /dev/null +++ b/xcape/Indices_calc.py @@ -0,0 +1,145 @@ + +def Indices_calc(p_2d, t_2d, td_2d, u_2d, v_2d, + p_s, t_s, td_s, u_s, v_s, + pres_lev_pos, aglh0, type_grid): + +# def Indices_calc(p_2d, t_2d, td_2d, u_2d, v_2d, q_2d, +# p_s, t_s, td_s, u_s, v_s, q_s, +# MUlevs , camu_s, cas_s, cin_s, caml_s, srh_rm1_s, srh_lm1_s, srh_rm3_s, srh_lm3_s, +# pres_lev_pos, aglh0, type_grid): + + import numpy as np + # nlev has to be the first dimension + # nlev here is the number of levels in 3d variables (without surface level) + nlev, ngrid = t_2d.shape + + # if above ground level height is 1 value + # (i.e. 2m or 10m) generate ngrid-long array with aglh0 + if np.isscalar(aglh0): + aglh_in = np.ones(ngrid)*aglh0 + else: + # add check on shape + aglh_in = aglh0 + + # type_grid type of vertical grid: 1 for model levels, 2 for pressure levels: + if type_grid == 1: + from xcape import stdheight_2D_model_lev + H2D, H_s = stdheight_2D_model_lev.loop_stdheight_ml(p_2d, t_2d, td_2d, + p_s, t_s, td_s, + aglh_in, + nlev, ngrid) + from xcape import Interp_model_lev + T2km = Interp_model_lev.interp_loop_ml(t_2d, H2D, t_s, H_s, 2000) + T3km = Interp_model_lev.interp_loop_ml(t_2d, H2D, t_s, H_s, 3000) + T4km = Interp_model_lev.interp_loop_ml(t_2d, H2D, t_s, H_s, 4000) + + T700 = Interp_model_lev.interp_loop_ml(t_2d, p_2d, t_s, p_s, 700) + T500 = Interp_model_lev.interp_loop_ml(t_2d, p_2d, t_s, p_s, 500) + z700 = Interp_model_lev.interp_loop_ml(H2D, p_2d, H_s, p_s, 700) + z500 = Interp_model_lev.interp_loop_ml(H2D, p_2d, H_s, p_s, 500) + + + FZL = Interp_model_lev.interp_loop_ml(H2D, t_2d, H_s, t_s, -0.001) + HT30 = Interp_model_lev.interp_loop_ml(H2D, t_2d, H_s, t_s, -30) + HT10 = Interp_model_lev.interp_loop_ml(H2D, t_2d, H_s, t_s, -10) + + U6km = Interp_model_lev.interp_loop_ml(u_2d, H2D, u_s, H_s, 6000) + V6km = Interp_model_lev.interp_loop_ml(v_2d, H2D, v_s, H_s, 6000) + U1km = Interp_model_lev.interp_loop_ml(u_2d, H2D, u_s, H_s, 1000) + V1km = Interp_model_lev.interp_loop_ml(v_2d, H2D, v_s, H_s, 1000) + + + + elif type_grid == 2: + from xcape import stdheight_2D_pressure_lev +# if flag_1d ==0: +# H2D, H_s = stdheight_2D_pressure_lev.loop_stdheight_pl(p_2d, t_2d, td_2d, +# p_s, t_s, td_s, +# aglh_in,pres_lev_pos, +# nlev, ngrid) +# if flag_1d ==1: + H2D, H_s = stdheight_2D_pressure_lev.loop_stdheight_pl1d(t_2d, td_2d, p_2d, + p_s, t_s, td_s, + aglh_in,pres_lev_pos, + nlev, ngrid) + from xcape import Interp_pressure_lev + T2km = Interp_pressure_lev.interp_loop_pl(t_2d, H2D, t_s, H_s, 2000,pres_lev_pos) + T3km = Interp_pressure_lev.interp_loop_pl(t_2d, H2D, t_s, H_s, 3000,pres_lev_pos) + T4km = Interp_pressure_lev.interp_loop_pl(t_2d, H2D, t_s, H_s, 4000,pres_lev_pos) + + T700 = Interp_pressure_lev.interp_loop_pl_y1d(t_2d, p_2d, t_s, p_s, 700,pres_lev_pos) + T500 = Interp_pressure_lev.interp_loop_pl_y1d(t_2d, p_2d, t_s, p_s, 500,pres_lev_pos) + z700 = Interp_pressure_lev.interp_loop_pl_y1d(H2D, p_2d, H_s, p_s, 700,pres_lev_pos) + z500 = Interp_pressure_lev.interp_loop_pl_y1d(H2D, p_2d, H_s, p_s, 500,pres_lev_pos) + + + FZL = Interp_pressure_lev.interp_loop_pl(H2D, t_2d, H_s, t_s, -0.00001,pres_lev_pos) + HT30 = Interp_pressure_lev.interp_loop_pl(H2D, t_2d, H_s, t_s, -30,pres_lev_pos) + HT10 = Interp_pressure_lev.interp_loop_pl(H2D, t_2d, H_s, t_s, -10,pres_lev_pos) + + U6km = Interp_pressure_lev.interp_loop_pl(u_2d, H2D, u_s, H_s, 6000,pres_lev_pos) + V6km = Interp_pressure_lev.interp_loop_pl(v_2d, H2D, v_s, H_s, 6000,pres_lev_pos) + U1km = Interp_pressure_lev.interp_loop_pl(u_2d, H2D, u_s, H_s, 1000,pres_lev_pos) + V1km = Interp_pressure_lev.interp_loop_pl(v_2d, H2D, v_s, H_s, 1000,pres_lev_pos) + + # Calculate LAPSE RATES - (Tsuperior-Tinferior)/(zsuperior-zinferior) + LAPSE24 = - (T4km-T2km)/2 + LAPSE3 = - (T3km-t_s)/3 + LAPSE700_500 = - (T500-T700)/((z500-z700)/1000) + THGZ = HT30-HT10 + # CALCULATE SBLCL + SBLCL = 125.* ( (t_s-t_2d[0])/2. -td_s) + # S06 + S06 = ( (U6km-u_s)**2 + (V6km-v_s)**2 )**0.5 + # S01 + S01 = ( (U1km-u_s)**2 + (V1km-v_s)**2 )**0.5 + + return LAPSE24, LAPSE3, LAPSE700_500, THGZ, S06, S01, SBLCL, T500, FZL + +# ################ +# ##### SHIP ##### +# q_mulev = np.where(MUlevs==1, q_s, q_2d[MUlevs-1,np.arange(q_2d.shape[1])]) + +# #SHIP +# SHIP = (camu_s* q_mulev* LAPSE700_500 * (-T500) * S06)/(44_000_000) +# SHIP = np.where(camu_s<1300, SHIP*(camu_s/1300), SHIP) +# SHIP = np.where(LAPSE700_500<5.8, SHIP*(LAPSE700_500/5.8), SHIP) +# SHIP = np.where(FZL<2400, SHIP*(FZL/2400), SHIP) + +# ################ +# ##### STP +# Aterm = cas_s/1500. +# Bterm = (2000.-SBLCL)/1000. +# Bterm[SBLCL<1000] = 1 +# Bterm[SBLCL>2000] = 0 +# Cterm_rm = srh_rm1_s/150. +# Cterm_lm = srh_lm1_s/150. +# Dterm = S06/20. +# Dterm[S06>30] = 1.5 +# Dterm[S06 < 12.5] = 0. +# Eterm = np.fabs(cin_s) +# Eterm[np.fabs(cin_s)>125]=0. +# Eterm[np.fabs(cin_s)<=125]=1. +# STP_rm = Aterm * Bterm * Cterm_rm * Dterm * Eterm +# STP_lm = Aterm * Bterm * Cterm_lm * Dterm * Eterm +# print(STP.shape) + +# ################ +# ##### SCP# +# scp1 = cas_s/1000. +# scp2_rm = srh_rm3_s/100. +# scp2_lm = srh_lm3_s/100. +# scp3 = S06/20. +# SCP_rm = scp1*scp2_rm*scp3*Eterm +# SCP_lm = scp1*scp2_lm*scp3*Eterm + +# ################ +# ##### EHI +# EHI3_rm = ((caml_s)*(srh_rm3_s))/(1.6*10**5) +# EHI1_rm = ((caml_s)*(srh_rm1_s))/(1.6*10**5) +# EHI3_lm = ((caml_s)*(srh_lm3_s))/(1.6*10**5) +# EHI1_lm = ((caml_s)*(srh_lm1_s))/(1.6*10**5) + + +# return LAPSE24, LAPSE3, LAPSE700_500, THGZ, S06, SHIP, STP_rm, STP_lm, SCP_rm, SCP_lm, EHI1_rm, EHI1_lm, EHI3_rm, EHI3_lm + diff --git a/xcape/Interp_model_lev.f90 b/xcape/Interp_model_lev.f90 new file mode 100644 index 0000000..299c634 --- /dev/null +++ b/xcape/Interp_model_lev.f90 @@ -0,0 +1,100 @@ +!----------------------------------------------------------------------- +!ccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc +!----------------------------------------------------------------------- + subroutine interp_loop_ml(INPUT, Y_TO_USE, INPUTs, Y_TO_USEs, LOC_in, nk, n2, OUTPUT) +!----------------------------------------------------------------------- +! interp_loop_ml - loop along n2 dimension to calculate +! interpolated INPUT onto Y_TO_USE for a given LOC_IN +! +! Author: Chiara Lepore @chiaral +! +! Disclaimer: This code is made available WITHOUT WARRANTY. +!----------------------------------------------------------------------- +! INPUT, Y_TO_USE, INPUTs, Y_TO_USEs = 3D and surface fields +! LOC_in = . value of Y_TO_USE_all to interpolate INPUT_all to +! nk = number of pressure levels +! n2 = number of grid point for which calculate interpolated value +!----------------------------------------------------------------------- + implicit none + + !f2py threadsafe + !f2py intent(out) :: OUTPUT + + integer, intent(in) :: nk,n2 + integer :: i, nk_all + real, dimension(nk,n2), intent(in) :: INPUT, Y_TO_USE + real, dimension(n2), intent(in) :: INPUTs, Y_TO_USEs + real, intent(in) :: LOC_in + + real, dimension(n2), intent(out) :: OUTPUT + real, dimension(nk+1) :: INPUT_all, Y_TO_USE_all + + nk_all = nk + 1 + do i = 1, n2 + INPUT_all(1) = INPUTs(i) + Y_TO_USE_all(1) = Y_TO_USEs(i) + INPUT_all(2:nk_all) = INPUT(:,i) + Y_TO_USE_all(2:nk_all) = Y_TO_USE(:,i) + + call DINTERP_DZ(INPUT_all,Y_TO_USE_all,LOC_in,OUTPUT(i),1,nk_all) + + enddo + return + end subroutine interp_loop_ml + + + + +!----------------------------------------------------------------------- +!ccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc +!----------------------------------------------------------------------- + + SUBROUTINE DINTERP_DZ(V3D,Z,LOC,V2D,N2,NZ) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(out) :: V2D + + INTEGER N2,NZ + REAL, INTENT(IN) :: V3D(NZ,N2) + REAL, INTENT(OUT) :: V2D(N2) + REAL, INTENT(IN) :: Z(NZ,N2) + REAL, INTENT(IN) :: LOC + + REAL VMSG + INTEGER J,KP,IP,IM + LOGICAL INTERP + REAL HEIGHT,W1,W2 + + HEIGHT = LOC + VMSG = -999999. + ! does vertical coordinate increase or decrease with increasing k? + ! set offset appropriately + + IP = 0 + IM = 1 + IF (Z(1,1).GT.Z(NZ,1)) THEN + IP = 1 + IM = 0 + END IF + !print *,' V3D,Z,LOC,V2D,N2,NZ in interp = ',V3D,Z,LOC,V2D,N2,NZ + + DO J = 1,N2 + ! Initialize to missing. Was initially hard-coded to -999999. + V2D(J) = VMSG + INTERP = .false. + KP = NZ + DO WHILE ((.NOT.INTERP) .AND. (KP.GE.2)) + IF (((Z(KP-IM,J).LE.HEIGHT).AND. (Z(KP-IP,J).GT.HEIGHT))) THEN + W2 = (HEIGHT-Z(KP-IM,J))/(Z(KP-IP,J)-Z(KP-IM,J)) + W1 = 1.D0 - W2 + V2D(J) = W1*V3D(KP-IM,J) + W2*V3D(KP-IP,J) + INTERP = .true. + ENDIF + KP = KP - 1 + ENDDO + ENDDO + !print *,' V2D in interp = ',V2D + + RETURN + END SUBROUTINE DINTERP_DZ diff --git a/xcape/Interp_model_lev.pyf b/xcape/Interp_model_lev.pyf new file mode 100644 index 0000000..dc40da4 --- /dev/null +++ b/xcape/Interp_model_lev.pyf @@ -0,0 +1,30 @@ +! -*- f90 -*- +! Note: the context of this file is case sensitive. + +python module Interp_model_lev ! in + interface ! in :Interp_model_lev + subroutine interp_loop_ml(input,y_to_use,inputs,y_to_uses,loc_in,nk,n2,output) ! in :Interp_model_lev:Interp_model_lev.f90 + threadsafe + real dimension(nk,n2),intent(in) :: input + real dimension(nk,n2),intent(in),depend(nk,n2) :: y_to_use + real dimension(n2),intent(in),depend(n2) :: inputs + real dimension(n2),intent(in),depend(n2) :: y_to_uses + real intent(in) :: loc_in + integer, optional,intent(in),check(shape(input,0)==nk),depend(input) :: nk=shape(input,0) + integer, optional,intent(in),check(shape(input,1)==n2),depend(input) :: n2=shape(input,1) + real dimension(n2),intent(out),depend(n2) :: output + end subroutine interp_loop_ml + subroutine dinterp_dz(v3d,z,loc,v2d,n2,nz) ! in :Interp_model_lev:Interp_model_lev.f90 + threadsafe + real dimension(nz,n2),intent(in) :: v3d + real dimension(nz,n2),intent(in),depend(nz,n2) :: z + real intent(in) :: loc + real dimension(n2),intent(out),depend(n2) :: v2d + integer, optional,check(shape(v3d,1)==n2),depend(v3d) :: n2=shape(v3d,1) + integer, optional,check(shape(v3d,0)==nz),depend(v3d) :: nz=shape(v3d,0) + end subroutine dinterp_dz + end interface +end python module Interp_model_lev + +! This file was auto-generated with f2py (version:2). +! See http://cens.ioc.ee/projects/f2py2e/ diff --git a/xcape/Interp_pressure_lev.f90 b/xcape/Interp_pressure_lev.f90 new file mode 100644 index 0000000..e186c7c --- /dev/null +++ b/xcape/Interp_pressure_lev.f90 @@ -0,0 +1,152 @@ +!----------------------------------------------------------------------- +!ccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc +!----------------------------------------------------------------------- + subroutine interp_loop_pl(INPUT, Y_TO_USE, INPUTs, Y_TO_USEs, LOC_in, start_3d, nk, n2, OUTPUT) +!----------------------------------------------------------------------- +! bunkers_loop - loop along n2 dimension to calculate RM,LM,Mean6kmwind. +! +! Author: Chiara Lepore @chiaral - John Allen @xebadir +! +! Disclaimer: This code is made available WITHOUT WARRANTY. +!----------------------------------------------------------------------- +! U3d,V3d = zonal and meridional wind +! AGLH3d = above ground level height +! nk = number of pressure levels +! n2 = number of grid point for which calculate RM,LM,Mean6kmwind +!----------------------------------------------------------------------- + implicit none + + !f2py threadsafe + !f2py intent(out) :: OUTPUT + + integer, intent(in) :: nk,n2 + integer :: i, nk_all, nk_pl_in, nk_start + + real, dimension(nk,n2), intent(in) :: INPUT, Y_TO_USE + real, dimension(n2), intent(in) :: INPUTs, Y_TO_USEs + real, intent(in) :: LOC_in + integer, dimension(n2), intent(in) :: start_3d + + real, dimension(n2), intent(out) :: OUTPUT + real, dimension(nk+1) :: INPUT_all, Y_TO_USE_all + + do i = 1, n2 + nk_start = start_3d(i) + ! compute number of used levels in 3d + nk_pl_in = nk - nk_start + 1 + ! compute number of used levels in 3d + surface + nk_all = nk_pl_in+1 + + INPUT_all(1) = INPUTs(i) + Y_TO_USE_all(1) = Y_TO_USEs(i) + INPUT_all(2:nk_all) = INPUT(nk_start:nk,i) + Y_TO_USE_all(2:nk_all) = Y_TO_USE(nk_start:nk,i) + + call DINTERP_DZ(INPUT_all,Y_TO_USE_all,LOC_in,OUTPUT(i),1,nk_all) + enddo + return + end subroutine Interp_loop_pl + +!----------------------------------------------------------------------- +!ccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc +!----------------------------------------------------------------------- + subroutine interp_loop_pl_Y1d(INPUT, Y_TO_USE, INPUTs, Y_TO_USEs, LOC_in, start_3d, nk, n2, OUTPUT) +!----------------------------------------------------------------------- +! bunkers_loop - loop along n2 dimension to calculate RM,LM,Mean6kmwind. +! +! Author: Chiara Lepore @chiaral - John Allen @xebadir +! +! Disclaimer: This code is made available WITHOUT WARRANTY. +!----------------------------------------------------------------------- +! U3d,V3d = zonal and meridional wind +! AGLH3d = above ground level height +! nk = number of pressure levels +! n2 = number of grid point for which calculate RM,LM,Mean6kmwind +!----------------------------------------------------------------------- + implicit none + + !f2py threadsafe + !f2py intent(out) :: OUTPUT + + integer, intent(in) :: nk,n2 + integer :: i, nk_all, nk_pl_in, nk_start + + real, dimension(nk,n2), intent(in) :: INPUT + real, dimension(nk,1), intent(in) :: Y_TO_USE + real, dimension(n2), intent(in) :: INPUTs, Y_TO_USEs + real, intent(in) :: LOC_in + integer, dimension(n2), intent(in) :: start_3d + + real, dimension(n2), intent(out) :: OUTPUT + real, dimension(nk+1) :: INPUT_all, Y_TO_USE_all + + do i = 1, n2 + nk_start = start_3d(i) + ! compute number of used levels in 3d + nk_pl_in = nk - nk_start + 1 + ! compute number of used levels in 3d + surface + nk_all = nk_pl_in+1 + + INPUT_all(1) = INPUTs(i) + Y_TO_USE_all(1) = Y_TO_USEs(i) + INPUT_all(2:nk_all) = INPUT(nk_start:nk,i) + Y_TO_USE_all(2:nk_all) = Y_TO_USE(nk_start:nk,1) + + call DINTERP_DZ(INPUT_all,Y_TO_USE_all,LOC_in,OUTPUT(i),1,nk_all) + enddo + return + end subroutine interp_loop_pl_Y1d + +!----------------------------------------------------------------------- +!ccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc +!----------------------------------------------------------------------- + + SUBROUTINE DINTERP_DZ(V3D,Z,LOC,V2D,N2,NZ) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(out) :: V2D + + INTEGER N2,NZ + REAL, INTENT(IN) :: V3D(NZ,N2) + REAL, INTENT(OUT) :: V2D(N2) + REAL, INTENT(IN) :: Z(NZ,N2) + REAL, INTENT(IN) :: LOC + + REAL VMSG + INTEGER J,KP,IP,IM + LOGICAL INTERP + REAL HEIGHT,W1,W2 + + HEIGHT = LOC + VMSG = -999999. + ! does vertical coordinate increase or decrease with increasing k? + ! set offset appropriately + + IP = 0 + IM = 1 + IF (Z(1,1).GT.Z(NZ,1)) THEN + IP = 1 + IM = 0 + END IF + !print *,' V3D,Z,LOC,V2D,N2,NZ in interp = ',V3D,Z,LOC,V2D,N2,NZ + + DO J = 1,N2 + ! Initialize to missing. Was initially hard-coded to -999999. + V2D(J) = VMSG + INTERP = .false. + KP = NZ + DO WHILE ((.NOT.INTERP) .AND. (KP.GE.2)) + IF (((Z(KP-IM,J).LE.HEIGHT).AND. (Z(KP-IP,J).GT.HEIGHT))) THEN + W2 = (HEIGHT-Z(KP-IM,J))/(Z(KP-IP,J)-Z(KP-IM,J)) + W1 = 1.D0 - W2 + V2D(J) = W1*V3D(KP-IM,J) + W2*V3D(KP-IP,J) + INTERP = .true. + ENDIF + KP = KP - 1 + ENDDO + ENDDO + !print *,' V2D in interp = ',V2D + + RETURN + END SUBROUTINE DINTERP_DZ diff --git a/xcape/Interp_pressure_lev.pyf b/xcape/Interp_pressure_lev.pyf new file mode 100644 index 0000000..bebe006 --- /dev/null +++ b/xcape/Interp_pressure_lev.pyf @@ -0,0 +1,43 @@ +! -*- f90 -*- +! Note: the context of this file is case sensitive. + +python module Interp_pressure_lev ! in + interface ! in :Interp_pressure_lev + subroutine interp_loop_pl(input,y_to_use,inputs,y_to_uses,loc_in,start_3d,nk,n2,output) ! in :Interp_pressure_lev:Interp_pressure_lev.f90 + threadsafe + real dimension(nk,n2),intent(in) :: input + real dimension(nk,n2),intent(in),depend(nk,n2) :: y_to_use + real dimension(n2),intent(in),depend(n2) :: inputs + real dimension(n2),intent(in),depend(n2) :: y_to_uses + real intent(in) :: loc_in + integer dimension(n2),intent(in),depend(n2) :: start_3d + integer, optional,intent(in),check(shape(input,0)==nk),depend(input) :: nk=shape(input,0) + integer, optional,intent(in),check(shape(input,1)==n2),depend(input) :: n2=shape(input,1) + real dimension(n2),intent(out),depend(n2) :: output + end subroutine interp_loop_pl + subroutine interp_loop_pl_y1d(input,y_to_use,inputs,y_to_uses,loc_in,start_3d,nk,n2,output) ! in :Interp_pressure_lev:Interp_pressure_lev.f90 + threadsafe + real dimension(nk,n2),intent(in) :: input + real dimension(nk,1),intent(in),depend(nk) :: y_to_use + real dimension(n2),intent(in),depend(n2) :: inputs + real dimension(n2),intent(in),depend(n2) :: y_to_uses + real intent(in) :: loc_in + integer dimension(n2),intent(in),depend(n2) :: start_3d + integer, optional,intent(in),check(shape(input,0)==nk),depend(input) :: nk=shape(input,0) + integer, optional,intent(in),check(shape(input,1)==n2),depend(input) :: n2=shape(input,1) + real dimension(n2),intent(out),depend(n2) :: output + end subroutine interp_loop_pl_y1d + subroutine dinterp_dz(v3d,z,loc,v2d,n2,nz) ! in :Interp_pressure_lev:Interp_pressure_lev.f90 + threadsafe + real dimension(nz,n2),intent(in) :: v3d + real dimension(nz,n2),intent(in),depend(nz,n2) :: z + real intent(in) :: loc + real dimension(n2),intent(out),depend(n2) :: v2d + integer, optional,check(shape(v3d,1)==n2),depend(v3d) :: n2=shape(v3d,1) + integer, optional,check(shape(v3d,0)==nz),depend(v3d) :: nz=shape(v3d,0) + end subroutine dinterp_dz + end interface +end python module Interp_pressure_lev + +! This file was auto-generated with f2py (version:2). +! See http://cens.ioc.ee/projects/f2py2e/ diff --git a/xcape/core.py b/xcape/core.py index 8939992..a601f1c 100644 --- a/xcape/core.py +++ b/xcape/core.py @@ -12,6 +12,7 @@ from .cape_numba import cape as _cape_numba from .srh import srh as _srh from .stdheight import stdheight as _stdheight +from .Indices_calc import Indices_calc as _Indices_calc def _prod(v): return reduce(lambda x, y: x*y, v) @@ -346,3 +347,108 @@ def _calc_srh_numpy(*args, srh_rm, srh_lm = _reshape_outputs(srh_2d_rm, srh_2d_lm, shape=original_shape) rm, lm, mean_6km = _reshape_outputs_uv_components(rm_2d, lm_2d, mean_6km_2d, shape=original_shape) return srh_rm, srh_lm, rm, lm, mean_6km + +def calc_indices(*args, **kwargs): + """ + Calculate cape for a set of profiles over the first axis of the arrays. + + Parameters + ---------- + p : array-like + Pressure in mb. + When vertical_lev='model', p.shape = t.shape = (nlev, x, y, ...) + When vertical_lev='pressure', p.shape = t.shape[0] = (nlev) + t : array-like + Temperature in Celsius + td : array-like + Dew point temperature in Celsius + u + v + + ps : array-like + Surface Pressure in mb. + ts : array-like + Surface Temperature in Celsius + tds : array-like + Surface Dew point temperature in Celsius + us, + vs + + vertical_lev : {'sigma', 'pressure'} + Which vertical grid is used + Returns + ------- + LAPSE24, LAPSE3, LAPSE700_500, THGZ, S06, SBLCL, T500, FZL : array-like + """ + if _any_dask_array(*args): + return _calc_indices_gufunc(*args, **kwargs) + else: + return _calc_indices_numpy(*args, **kwargs) + +def _calc_indices_gufunc(*args, **kwargs): + + + signatureA = "(i),(i),(i),(i),(i),(),(),(),(),()" + signatureB = "->(),()" + output_dtypes = ('f4','f4','f4','f4','f4','f4','f4','f4') + if (kwargs['vertical_lev']=='pressure'): + #ADD 1 INPUT FOR PRES_LEV_POS + signatureA = signatureA + ",()" + signature = signatureA + signatureB + return da.apply_gufunc(_calc_indices_numpy, signature, + *args, + output_dtypes=output_dtypes, + axis=-1, + vectorize=False, + **kwargs) + + +# the numpy version of the algorithm +def _calc_indices_numpy(*args, + vertical_lev='sigma'): + + #21 + p, t, td, u, v, ps, ts, tds, us, vs = args + + original_shape = t.shape #shape of 3D variable, i.e. p + original_surface_shape = ts.shape #shape of surface variable, i.e. ps + + # after this, all arrays are 2d shape (nlevs, npoints) + p_s1d, t_s1d, td_s1d, u_s1d, v_s1d = _reshape_surface_inputs(ps, ts, tds, us, vs) + + if len(p.shape) == 1: + t_2d, td_2d, u_2d, v_2d = _reshape_inputs(t, td, u, v) + p_2d = _reshape_inputs(p)[0] + flag_1d = 1 + # calculate pres_lev_pos + temp_index = (p_s1d-p_2d) + pres_lev_pos = np.ma.masked_less(temp_index,0).argmin(axis=0) + # fortran convention + pres_lev_pos = pres_lev_pos+1 + + elif (p.shape == t.shape)&(vertical_lev=='sigma'): + p_2d, t_2d, td_2d, u_2d, v_2d = _reshape_inputs(p, t, td, u, v) + flag_1d = 0 + pres_lev_pos = 1 + elif (p.shape == t.shape)&(vertical_lev=='pressure'): + raise ValueError("P should be 1d") + + + _vertical_lev_options_ ={'sigma':1, 'pressure':2} + kwargs = dict(type_grid=_vertical_lev_options_[vertical_lev]) + + + LAPSE24_1d, LAPSE3_1d, LAPSE700_500_1d, THGZ_1d, S06_1d, S01_1d, SBLCL_1d, T500_1d, FZL_1d = _Indices_calc( + p_2d, t_2d, td_2d, u_2d, v_2d, + p_s1d, t_s1d, td_s1d, u_s1d, v_s1d, + pres_lev_pos, aglh0 = 2., + **kwargs) + + #_reshape_outputs returns a list + LAPSE24, LAPSE3, LAPSE700_500, THGZ, S06, S01, SBLCL, T500, FZL = _reshape_outputs(LAPSE24_1d, LAPSE3_1d, + LAPSE700_500_1d, THGZ_1d, + S06_1d, S01_1d, SBLCL_1d, T500_1d, FZL_1d, + shape=original_shape) + + return LAPSE24, LAPSE3, LAPSE700_500, THGZ, S06, S01, SBLCL, T500, FZL +