diff --git a/fortran/constants.f90 b/fortran/constants.f90 new file mode 100644 index 0000000..c74985c --- /dev/null +++ b/fortran/constants.f90 @@ -0,0 +1,10 @@ +MODULE constants + INTEGER :: ERRLEN=512 + REAL(KIND=8), PARAMETER :: P1000MB=100000.D0 + REAL(KIND=8), PARAMETER :: R_D=287.D0 + REAL(KIND=8), PARAMETER :: CP=7.D0*R_D/2.D0 + REAL(KIND=8), PARAMETER :: R=287.04D0 + REAL(KIND=8), PARAMETER :: G=9.81D0 + REAL(KIND=8), PARAMETER :: GAMMA=0.0065D0 +END MODULE constants + diff --git a/fortran/wrf_user.f90 b/fortran/wrf_user.f90 new file mode 100644 index 0000000..7de0a98 --- /dev/null +++ b/fortran/wrf_user.f90 @@ -0,0 +1,908 @@ +! NCLFORTSTART +SUBROUTINE DCOMPUTEPI(pi, pressure, nx, ny, nz) + USE constants, ONLY : P1000MB, R_D, CP + + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: pi + + INTEGER, INTENT(IN) :: nx, ny, nz + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(OUT) :: pi + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: pressure +! NCLEND + + INTEGER i, j, k + + !REAL(KIND=8), PARAMETER :: P1000MB=100000.D0, R_D=287.D0, CP=7.D0*R_D/2.D0 + + DO k = 1,nz + DO j = 1,ny + DO i = 1,nx + pi(i,j,k) = (pressure(i,j,k)/P1000MB)**(R_D/CP) + END DO + END DO + END DO + +END SUBROUTINE DCOMPUTEPI + +! Temperature from potential temperature in kelvin. +!NCLFORTSTART +SUBROUTINE DCOMPUTETK(tk, pressure, theta, nx) + USE constants, ONLY : P1000MB, R_D, CP + + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: tk + + INTEGER, INTENT(IN) :: nx + REAL(KIND=8) :: pi + REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: pressure + REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: theta + REAL(KIND=8), DIMENSION(nx), INTENT(OUT) :: tk + +! NCLEND + + INTEGER :: i + + !REAL(KIND=8), PARAMETER :: P1000MB=100000.D0, R_D=287.D0, CP=7.D0*R_D/2.D0 + + DO i = 1,nx + pi = (pressure(i)/P1000MB)**(R_D/CP) + tk(i) = pi*theta(i) + END DO + + RETURN + +END SUBROUTINE DCOMPUTETK + + +! NCLFORTSTART +SUBROUTINE DINTERP3DZ(data3d, out2d, zdata, desiredloc, nx, ny, nz, missingval) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: out2d + + INTEGER, INTENT(IN) :: nx,ny,nz + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: data3d + REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: out2d + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: zdata + REAL(KIND=8), INTENT(IN) :: desiredloc + REAL(KIND=8), INTENT(IN) :: missingval + +! NCLEND + + INTEGER :: i,j,kp,ip,im + LOGICAL :: dointerp + REAL(KIND=8) :: height,w1,w2 + + height = desiredloc + + ! does vertical coordinate increase or decrease with increasing k? + ! set offset appropriately + + ip = 0 + im = 1 + IF (zdata(1,1,1) .GT. zdata(1,1,nz)) THEN + ip = 1 + im = 0 + END IF + + DO i = 1,nx + DO j = 1,ny + ! Initialize to missing. Was initially hard-coded to -999999. + out2d(i,j) = missingval + dointerp = .FALSE. + kp = nz + + DO WHILE ((.NOT. dointerp) .AND. (kp >= 2)) + IF (((zdata(i,j,kp-im) < height) .AND. (zdata(i,j,kp-ip) > height))) THEN + w2 = (height-zdata(i,j,kp-im))/(zdata(i,j,kp-ip)-zdata(i,j,kp-im)) + w1 = 1.D0 - w2 + out2d(i,j) = w1*data3d(i,j,kp-im) + w2*data3d(i,j,kp-ip) + dointerp = .TRUE. + END IF + kp = kp - 1 + END DO + + END DO + END DO + + RETURN + +END SUBROUTINE DINTERP3DZ + +! PORT DZSTAG HERE + +! NCLFORTSTART +SUBROUTINE DZSTAG(znew, nx, ny, nz, z, nxz, nyz ,nzz, terrain) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: znew + + INTEGER, INTENT(IN) :: nx, ny, nz, nxz, nyz, nzz + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(OUT) :: znew + REAL(KIND=8), DIMENSION(nxz,nyz,nzz), INTENT(IN) :: z + REAL(KIND=8), DIMENSION(nxz,nyz), INTENT(IN) :: terrain +! NCLEND + + INTEGER :: i,j,k,ii,im1,jj,jm1 + + ! check for u, v, or w (x,y,or z) staggering + ! for x and y stag, avg z to x, y, point + IF (nx .GT. nxz) THEN + DO k = 1,nz + DO j = 1,ny + DO i = 1,nx + ii = MIN0(i,nxz) + im1 = MAX0(i-1,1) + znew(i,j,k) = 0.5D0 * (z(ii,j,k) + z(im1,j,k)) + END DO + END DO + END DO + + ELSE IF (ny .GT. nyz) THEN + DO k = 1,nz + DO j = 1,NY + jj = MIN0(j,nyz) + jm1 = MAX0(j-1,1) + DO i = 1,nx + znew(i,j,k) = 0.5D0 * (z(i,jj,k) + z(i,jm1,k)) + END DO + END DO + END DO + + ! w (z) staggering + ELSE IF (nz .GT. nzz) THEN + DO j = 1,ny + DO i = 1,nx + znew(i,j,1) = terrain(i,j) + END DO + END DO + + DO k = 2,nz + DO j = 1,ny + DO i = 1,nx + znew(i,j,k) = znew(i,j,k-1) + 2.D0 * (z(i,j,k-1) - znew(i,j,k-1)) + END DO + END DO + END DO + + END IF + + RETURN + + END SUBROUTINE DZSTAG + +! NCLFORTSTART +SUBROUTINE DINTERP2DXY(v3d, v2d, xy, nx, ny, nz, nxy) + + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: v2d + + INTEGER, INTENT(IN) :: nx, ny, nz, nxy + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: v3d + REAL(KIND=8), DIMENSION(nxy,nz), INTENT(OUT) :: v2d + REAL(KIND=8), DIMENSION(2,nxy), INTENT(IN) :: xy + +! NCLEND + + INTEGER :: i, j, k, ij + REAL(KIND=8) :: w11, w12, w21, w22, wx, wy + + DO ij = 1,nxy + i = MAX0(1,MIN0(nx-1,INT(xy(1,ij)+1))) + j = MAX0(1,MIN0(ny-1,INT(xy(2,ij)+1))) + wx = DBLE(i+1) - (xy(1,ij)+1) + wy = DBLE(j+1) - (xy(2,ij)+1) + w11 = wx*wy + w21 = (1.D0-wx)*wy + w12 = wx*(1.D0-wy) + w22 = (1.D0-wx)* (1.D0-wy) + DO k = 1,nz + v2d(ij,k) = w11*v3d(i,j,k) + w21*v3d(i+1,j,k) + & + w12*v3d(i,j+1,k) + w22*v3d(i+1,j+1,k) + END DO + END DO + + RETURN + +END SUBROUTINE DINTERP2DXY + + +! NCLFORTSTART +SUBROUTINE DINTERP1D(v_in, v_out, z_in, z_out, vmsg, nz_in, nz_out) + + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: v_out + + INTEGER, INTENT(IN) :: nz_in, nz_out + REAL(KIND=8), DIMENSION(nz_in), INTENT(IN) :: v_in, z_in + REAL(KIND=8), DIMENSION(nz_out), INTENT(IN) :: z_out + REAL(KIND=8), DIMENSION(nz_out), INTENT(OUT) :: v_out + REAL(KIND=8), INTENT(IN) :: vmsg + +! NCLEND + + INTEGER :: kp,k,im,ip + LOGICAL :: interp + REAL(KIND=8) :: height,w1,w2 + + ! does vertical coordinate increase of decrease with increasing k? + ! set offset appropriately + + ip = 0 + im = 1 + IF (z_in(1) .GT. z_in(nz_in)) THEN + ip = 1 + im = 0 + END IF + + DO k = 1,nz_out + v_out(k) = vmsg + + interp = .FALSE. + kp = nz_in + height = z_out(k) + + DO WHILE ((.NOT. interp) .AND. (kp .GE. 2)) + IF (((z_in(kp-im) .LE. height) .AND. (z_in(kp-ip) .GT. height))) THEN + w2 = (height - z_in(kp-im))/(z_in(kp-ip) - z_in(kp-im)) + w1 = 1.d0 - w2 + v_out(k) = w1*v_in(kp-im) + w2*v_in(kp-ip) + interp = .TRUE. + END IF + kp = kp - 1 + END DO + END DO + + RETURN + +END SUBROUTINE DINTERP1D + +! This routine assumes +! index order is (i,j,k) +! wrf staggering +! +! units: pressure (Pa), temperature(K), height (m), mixing ratio +! (kg kg{-1}) availability of 3d p, t, and qv; 2d terrain; 1d +! half-level zeta string +! output units of SLP are Pa, but you should divide that by 100 for the +! weather weenies. +! virtual effects are included +! + +! NCLFORTSTART +SUBROUTINE DCOMPUTESEAPRS(nx, ny, nz, z, t, p, q, sea_level_pressure, & + t_sea_level, t_surf, level, errstat, errmsg) + USE constants, ONLY : ERRLEN, R, G, GAMMA + + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: sea_level_pressure + + ! Estimate sea level pressure. + INTEGER, INTENT(IN) :: nx, ny, nz + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: z + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: t, p, q + ! The output is the 2d sea level pressure. + REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: sea_level_pressure + INTEGER, DIMENSION(nx,ny), INTENT(INOUT) :: level + REAL(KIND=8), DIMENSION(nx,ny), INTENT(INOUT) :: t_surf, t_sea_level +! NCLFORTEND + + + INTEGER, OPTIONAL, INTENT(INOUT) :: errstat + CHARACTER(LEN=ERRLEN), OPTIONAL, INTENT(INOUT) :: errmsg + + ! Some required physical constants: + + !REAL(KIND=8), PARAMETER :: R=287.04D0, G=9.81D0, GAMMA=0.0065D0 + + ! Specific constants for assumptions made in this routine: + REAL(KIND=8), PARAMETER :: TC=273.16D0+17.5D0, PCONST=10000 + + LOGICAL, PARAMETER :: ridiculous_mm5_test=.TRUE. + ! PARAMETER (ridiculous_mm5_test = .FALSE.) + + ! Local variables: + + INTEGER :: i, j, k + INTEGER :: klo, khi + + REAL(KIND=8) :: plo, phi, tlo, thi, zlo, zhi + REAL(KIND=8) :: p_at_pconst, t_at_pconst, z_at_pconst + REAL(KIND=8) :: z_half_lowest + + LOGICAL :: l1, l2, l3, found + + ! Find least zeta level that is PCONST Pa above the surface. We + ! later use this level to extrapolate a surface pressure and + ! temperature, which is supposed to reduce the effect of the diurnal + ! heating cycle in the pressure field. + + IF (PRESENT(errstat)) THEN + errstat = 0 + END IF + + DO j = 1,ny + DO i = 1,nx + level(i,j) = -1 + + k = 1 + found = .FALSE. + DO WHILE ((.NOT. found) .AND. (k <= nz)) + IF (p(i,j,k) < p(i,j,1)-PCONST) THEN + level(i,j) = k + found = .TRUE. + END IF + k = k + 1 + END DO + + IF (level(i,j) == -1) THEN + IF (PRESENT(errstat)) THEN + errstat = 1 + errmsg = 'Error in finding 100 hPa up' + RETURN + ELSE + PRINT '(A,I4,A)','Troubles finding level ', NINT(PCONST)/100,' above ground.' + PRINT '(A,I4,A,I4,A)','Problems first occur at (',I,',',J,')' + PRINT '(A,F6.1,A)','Surface pressure = ',p(i,j,1)/100,' hPa.' + STOP 'Error in finding 100 hPa up' + END IF + + END IF + END DO + END DO + + ! Get temperature PCONST Pa above surface. Use this to extrapolate + ! the temperature at the surface and down to sea level. + + DO J = 1,ny + DO I = 1,nx + + klo = MAX(level(i,j)-1,1) + khi = MIN(klo+1,nz-1) + + IF (klo == khi) THEN + IF (PRESENT(errstat)) THEN + errstat = 1 + errmsg = 'Error trapping levels' + RETURN + ELSE + PRINT '(A)','Trapping levels are weird.' + PRINT '(A,I3,A,I3,A)','klo = ',klo,', khi = ',khi,': and they should not be equal.' + STOP 'Error trapping levels' + END IF + END IF + + plo = p(i,j,klo) + phi = p(i,j,khi) + tlo = t(i,j,klo)* (1.D0+0.608D0*q(i,j,klo)) + thi = t(i,j,khi)* (1.D0+0.608D0*q(i,j,khi)) + ! zlo = zetahalf(klo)/ztop*(ztop-terrain(i,j))+terrain(i,j) + ! zhi = zetahalf(khi)/ztop*(ztop-terrain(i,j))+terrain(i,j) + zlo = z(i,j,klo) + zhi = z(i,j,khi) + + p_at_pconst = p(i,j,1) - PCONST + t_at_pconst = thi - (thi-tlo)*LOG(p_at_pconst/phi)*LOG(plo/phi) + z_at_pconst = zhi - (zhi-zlo)*LOG(p_at_pconst/phi)*LOG(plo/phi) + + t_surf(i,j) = t_at_pconst * (p(i,j,1)/p_at_pconst)**(GAMMA*R/G) + t_sea_level(i,j) = t_at_pconst + GAMMA*z_at_pconst + + END DO + END DO + + ! If we follow a traditional computation, there is a correction to the + ! sea level temperature if both the surface and sea level + ! temperatures are *too* hot. + + IF (ridiculous_mm5_test) THEN + DO J = 1,ny + DO I = 1,nx + l1 = t_sea_level(i,j) < TC + l2 = t_surf(i,j) <= TC + l3 = .NOT. l1 + IF (l2 .AND. l3) THEN + t_sea_level(i,j) = TC + ELSE + t_sea_level(i,j) = TC - 0.005D0*(t_surf(i,j)-TC)**2 + END IF + END DO + END DO + END IF + + ! The grand finale: ta da! + DO J = 1,ny + DO I = 1,nx + ! z_half_lowest=zetahalf(1)/ztop*(ztop-terrain(i,j))+terrain(i,j) + z_half_lowest = z(i,j,1) + + ! Convert to hPa in this step, by multiplying by 0.01. The original + ! Fortran routine didn't do this, but the NCL script that called it + ! did, so we moved it here. + sea_level_pressure(i,j) = 0.01 * (p(i,j,1)*EXP((2.D0*G*z_half_lowest)/& + (R*(t_sea_level(i,j)+t_surf(i,j))))) + END DO + END DO + + ! PRINT *,'sea pres input at weird location i=20,j=1,k=1' + ! PRINT *,'t=',t(20,1,1),t(20,2,1),t(20,3,1) + ! PRINT *,'z=',z(20,1,1),z(20,2,1),z(20,3,1) + ! PRINT *,'p=',p(20,1,1),p(20,2,1),p(20,3,1) + ! PRINT *,'slp=',sea_level_pressure(20,1), + ! * sea_level_pressure(20,2),sea_level_pressure(20,3) + + RETURN + +END SUBROUTINE DCOMPUTESEAPRS + + +! Double precision version. If you make a change here, you +! must make the same change below to filter2d. + +! NCLFORTSTART +SUBROUTINE DFILTER2D(a, b, nx, ny, it, missing) + + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: a + + INTEGER, INTENT(IN) :: nx, ny, it + REAL(KIND=8), DIMENSION(nx, ny), INTENT(INOUT) :: a + REAL(KIND=8), INTENT(IN) :: missing + REAL(KIND=8), DIMENSION(nx, ny), INTENT(INOUT) :: b + +! NCLEND + + REAL(KIND=8), PARAMETER :: COEF=0.25D0 + + INTEGER :: i, j, iter + + DO iter=1,it + DO j=1,ny + DO i = 1,nx + b(i,j) = a(i,j) + END DO + END DO + + DO j=2,ny-1 + DO i=1,nx + IF (b(i,j-1) .EQ. missing .OR. b(i,j) .EQ. missing .OR. & + b(i,j+1) .EQ. missing) THEN + a(i,j) = a(i,j) + ELSE + a(i,j) = a(i,j) + COEF*(b(i,j-1) - 2*b(i,j) + b(i,j+1)) + END IF + END DO + END DO + + DO j=1,ny + DO i=2,nx-1 + IF (b(i-1,j) .EQ. missing .OR. b(i,j) .EQ. missing .OR. & + b(i+1,j) .EQ. missing) THEN + a(i,j) = a(i,j) + ELSE + a(i,j) = a(i,j) + COEF*(b(i-1,j) - 2*b(i,j) + b(i+1,j)) + END IF + END DO + END DO + ! do j=1,ny + ! do i=1,nx + ! b(i,j) = a(i,j) + ! enddo + ! enddo + ! do j=2,ny-1 + ! do i=1,nx + ! a(i,j) = a(i,j) - .99*coef*(b(i,j-1)-2*b(i,j)+b(i,j+1)) + ! enddo + ! enddo + ! do j=1,ny + ! do i=2,nx-1 + ! a(i,j) = a(i,j) - .99*coef*(b(i-1,j)-2*b(i,j)+b(i+1,j)) + ! enddo + ! enddo + END DO + + RETURN + +END SUBROUTINE DFILTER2D + +! Single precision version. If you make a change here, you +! must make the same change below to dfilter2d. + +! NCLFORTSTART +SUBROUTINE FILTER2D(a, b, nx, ny, it, missing) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: a + + INTEGER, INTENT(IN) :: nx, ny, it + REAL(KIND=4), DIMENSION(nx, ny), INTENT(INOUT) :: a + REAL(KIND=4), INTENT(IN) :: missing + REAL(KIND=4), DIMENSION(nx, ny), INTENT(INOUT) :: b + +! NCLEND + + REAL(KIND=4), PARAMETER :: COEF=0.25 + + INTEGER :: i, j, iter + + DO iter=1,it + DO j=1,ny + DO i = 1,nx + b(i,j) = a(i,j) + END DO + END DO + + DO j=2,ny-1 + DO i=1,nx + IF (b(i,j-1) .EQ. missing .OR. b(i,j) .EQ. missing .OR. & + b(i,j+1) .EQ. missing) THEN + a(i,j) = a(i,j) + ELSE + a(i,j) = a(i,j) + COEF*(b(i,j-1)-2*b(i,j) + b(i,j+1)) + END IF + END DO + END DO + + DO j=1,ny + DO i=2,nx-1 + IF (b(i-1,j) .EQ. missing .OR. b(i,j) .EQ. missing .OR. & + b(i+1,j) .EQ. missing) THEN + a(i,j) = a(i,j) + ELSE + a(i,j) = a(i,j) + COEF*(b(i-1,j)-2*b(i,j)+b(i+1,j)) + END IF + END DO + END DO + ! do j=1,ny + ! do i=1,nx + ! b(i,j) = a(i,j) + ! enddo + ! enddo + ! do j=2,ny-1 + ! do i=1,nx + ! a(i,j) = a(i,j) - .99*coef*(b(i,j-1)-2*b(i,j)+b(i,j+1)) + ! enddo + ! enddo + ! do j=1,ny + ! do i=2,nx-1 + ! a(i,j) = a(i,j) - .99*coef*(b(i-1,j)-2*b(i,j)+b(i+1,j)) + ! enddo + ! enddo + END DO + + RETURN + +END SUBROUTINE FILTER2D + + +! NCLFORTSTART +SUBROUTINE DCOMPUTERH(qv, p, t, rh, nx) + + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: rh + + INTEGER, INTENT(IN) :: nx + REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: qv,p,t + REAL(KIND=8), DIMENSION(nx), INTENT(OUT) :: rh + +! NCLEND + + REAL(KIND=8), PARAMETER :: SVP1=0.6112D0,SVP2=17.67D0,SVP3=29.65D0,SVPT0=273.15D0 + + INTEGER :: i + REAL(KIND=8) :: qvs,es,pressure,temperature + REAL(KIND=8), PARAMETER :: R_D=287.D0,R_V=461.6D0,EP_2=R_D/R_V + REAL(KIND=8), PARAMETER :: EP_3=0.622D0 + + DO i = 1,nx + pressure = p(i) + temperature = t(i) + ! es = 1000.*svp1* + es = 10.D0*SVP1*EXP(SVP2* (temperature-SVPT0)/(temperature-SVP3)) + ! qvs = ep_2*es/(pressure-es) + qvs = EP_3*es/ (0.01D0*pressure- (1.D0-EP_3)*es) + ! rh = 100*amax1(1., qv(i)/qvs) + ! rh(i) = 100.*qv(i)/qvs + rh(i) = 100.D0*DMAX1(DMIN1(qv(i)/qvs,1.0D0),0.0D0) + END DO + + RETURN + +END SUBROUTINE DCOMPUTERH + + +! NCLFORTSTART +SUBROUTINE DGETIJLATLONG(lat_array, long_array, lat, longitude, ii, jj, nx, ny, imsg) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: ii,jj + + INTEGER, INTENT(IN) :: nx,ny,imsg + INTEGER, INTENT(OUT) :: ii,jj + REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: lat_array,long_array + REAL(KIND=8) :: lat,longitude +! NCLEND + + REAL(KIND=8) :: longd,latd + INTEGER :: i,j + REAL(KIND=8) :: ir,jr + REAL(KIND=8) :: dist_min,dist + + ! init to missing. was hard-coded to -999 initially. + ir = imsg + jr = imsg + + dist_min = 1.d+20 + DO j = 1,ny + DO i = 1,nx + latd = (lat_array(i,j)-lat)**2 + longd = (long_array(i,j)-longitude)**2 + ! longd = dmin1((long_array(i,j)-longitude)**2, & + ! (long_array(i,j)+longitude)**2) + dist = SQRT(latd+longd) + IF (dist_min .GT. dist) THEN + dist_min = dist + ir = DBLE(i) + jr = DBLE(j) + END IF + END DO + END DO + + ! The original version of this routine returned IR and JR. But, then + ! the NCL script that called this routine was converting IR and JR + ! to integer, so why not just return II and JJ? + + ! Also, I'm subtracing 1 here, because it will be returned to NCL + ! script which has 0-based indexing. + + IF (ir .NE. imsg .AND. jr .NE. imsg) THEN + ii = NINT(ir)-1 + jj = NINT(jr)-1 + ELSE + ii = imsg + jj = imsg + END IF + + ! we will just return the nearest point at present + + RETURN + +END SUBROUTINE DGETIJLATLONG + +! You need to modify the C-WRAPPER in NCL for this to work + +! NCLFORTSTART +SUBROUTINE DCOMPUTEUVMET(u, v, uvmet, longca,longcb,flong,flat, & + cen_long, cone, rpd, nx, ny, nz, nxp1, nyp1, & + istag, is_msg_val, umsg, vmsg, uvmetmsg) + IMPLICIT NONE + + ! ISTAG should be 0 if the U,V grids are not staggered. + ! That is, NY = NYP1 and NX = NXP1. + + !f2py threadsafe + !f2py intent(in,out) :: uvmet + + INTEGER,INTENT(IN) :: nx,ny,nz,nxp1,nyp1,istag + LOGICAL,INTENT(IN) :: is_msg_val + REAL(KIND=8), DIMENSION(nxp1,ny,nz), INTENT(IN):: u + REAL(KIND=8), DIMENSION(nx,nyp1,nz), INTENT(IN) :: v + REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: flong + REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: flat + REAL(KIND=8), DIMENSION(nx,ny), INTENT(INOUT) :: longca + REAL(KIND=8), DIMENSION(nx,ny), INTENT(INOUT) :: longcb + REAL(KIND=8), INTENT(IN) :: cen_long,cone,rpd + REAL(KIND=8), INTENT(IN) :: umsg,vmsg,uvmetmsg + REAL(KIND=8), DIMENSION(nx,ny,nz,2), INTENT(OUT) :: uvmet + + +! NCLEND + + INTEGER :: i,j,k + REAL(KIND=8) :: uk,vk + + ! msg stands for missing value in this code + ! WRITE (6,FMT=*) ' in compute_uvmet ',NX,NY,NXP1,NYP1,ISTAG + + DO j = 1,ny + DO i = 1,nx + + longca(i,j) = flong(i,j) - cen_long + IF (longca(i,j).GT.180.D0) THEN + longca(i,j) = longca(i,j) - 360.D0 + END IF + IF (longca(i,j).LT.-180.D0) THEN + longca(i,j) = longca(i,j) + 360.D0 + END IF + IF (flat(i,j).LT.0.D0) THEN + longcb(i,j) = -longca(i,j)*cone*rpd + ELSE + longcb(i,j) = longca(i,j)*cone*rpd + END IF + + longca(i,j) = COS(longcb(i,j)) + longcb(i,j) = SIN(longcb(i,j)) + + END DO + END DO + + ! WRITE (6,FMT=*) ' computing velocities ' + DO k = 1,nz + DO j = 1,ny + DO i = 1,nx + IF (istag.EQ.1) THEN + IF (is_msg_val .AND. (u(i,j,k) .EQ. umsg .OR. v(i,j,k) .EQ. vmsg & + .OR. u(i+1,j,k) .EQ. umsg .OR. v(i,j+1,k) .EQ. vmsg)) THEN + uvmet(i,j,k,1) = uvmetmsg + uvmet(i,j,k,2) = uvmetmsg + ELSE + uk = 0.5D0* (u(i,j,k)+u(i+1,j,k)) + vk = 0.5D0* (v(i,j,k)+v(i,j+1,k)) + uvmet(i,j,k,1) = vk*longcb(i,j) + uk*longca(i,j) + uvmet(i,j,k,2) = vk*longca(i,j) - uk*longcb(i,j) + END IF + ELSE + IF (is_msg_val .AND. (u(i,j,k) .EQ. umsg .OR. v(i,j,k) .EQ. vmsg)) THEN + uvmet(i,j,k,1) = uvmetmsg + uvmet(i,j,k,2) = uvmetmsg + ELSE + uk = u(i,j,k) + vk = v(i,j,k) + uvmet(i,j,k,1) = vk*longcb(i,j) + uk*longca(i,j) + uvmet(i,j,k,2) = vk*longca(i,j) - uk*longcb(i,j) + END IF + END IF + END DO + END DO + END DO + + RETURN + +END SUBROUTINE DCOMPUTEUVMET + + + + +! This was originally a routine that took 2D input arrays. Since +! the NCL C wrapper routine can handle multiple dimensions, it's +! not necessary to have anything bigger than 1D here. + +! NCLFORTSTART +SUBROUTINE DCOMPUTETD(td, pressure, qv_in, nx) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: td + + INTEGER, INTENT(IN) :: nx + REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: pressure + REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: qv_in + REAL(KIND=8), DIMENSION(nx), INTENT(OUT) :: td + +! NCLEND + + REAL(KIND=8) :: qv,tdc + + INTEGER :: i + + DO i = 1,nx + qv = DMAX1(qv_in(i),0.D0) + ! vapor pressure + tdc = qv*pressure(i)/ (.622D0 + qv) + + ! avoid problems near zero + tdc = DMAX1(tdc,0.001D0) + td(i) = (243.5D0*LOG(tdc)-440.8D0) / (19.48D0-LOG(tdc)) + END DO + + RETURN + +END SUBROUTINE DCOMPUTETD + +! NCLFORTSTART +SUBROUTINE DCOMPUTEICLW(iclw, pressure, qc_in, nx, ny, nz) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: iclw + + INTEGER, INTENT(IN) :: nx,ny,nz + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: pressure + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: qc_in + REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: iclw + +! NCLEND + + REAL(KIND=8) :: qclw,dp + REAL(KIND=8), PARAMETER :: GG = 1000.d0/9.8d0 + INTEGER i,j,k + + DO j = 1,ny + DO i = 1,nx + iclw(i,j) = 0.d0 + END DO + END DO + + DO j = 3,ny - 2 + DO i = 3,nx - 2 + DO k = 1,nz + qclw = DMAX1(qc_in(i,j,k),0.d0) + IF (k.EQ.1) THEN + dp = pressure(i,j,k-1) - pressure(i,j,k) + ELSE IF (k.EQ.nz) then + dp = pressure(i,j,k) - pressure(i,j,k+1) + ELSE + dp = (pressure(i,j,k-1) - pressure(i,j,k+1)) / 2.d0 + END IF + iclw(i,j) = iclw(i,j) + qclw*dp*GG + END DO + END DO + END DO + + RETURN + +END SUBROUTINE DCOMPUTEICLW + +SUBROUTINE testfunc(a, b, nx, ny, nz, errstat, errstr) + IMPLICIT NONE + + !f2py threadsafe + !f2py intent(in,out) :: b + + INTEGER, INTENT(IN) :: nx, ny, nz + + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: a + REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(OUT) :: b + INTEGER, INTENT(INOUT), OPTIONAL :: errstat + CHARACTER(LEN=512), INTENT(INOUT), OPTIONAL :: errstr + + INTEGER :: i,j,k + + IF (PRESENT(errstat)) THEN + errstat = 0 + errstr = 'test string worked' + END IF + + + DO k=1,nz + DO j=1,ny + DO i=1,nx + b(i,j,k) = a(i,j,k) + END DO + END DO + END DO + + IF (PRESENT(errstat)) THEN + errstat = 1 + END IF + + RETURN + +END SUBROUTINE testfunc + + + + + + + + + diff --git a/fortran/wrffortran.pyf b/fortran/wrffortran.pyf new file mode 100644 index 0000000..9e38207 --- /dev/null +++ b/fortran/wrffortran.pyf @@ -0,0 +1,183 @@ +! -*- f90 -*- +! Note: the context of this file is case sensitive. + +python module _wrffortran ! in + interface ! in :_wrffortran + module constants ! in :_wrffortran:constants.f90 + real(kind=8), parameter,optional :: p1000mb=100000.d0 + real(kind=8), parameter,optional :: r=287.04d0 + real(kind=8), parameter,optional :: g=9.81d0 + real(kind=8), parameter,optional :: r_d=287.d0 + real(kind=8), parameter,optional,depend(r_d) :: cp=7.d0*r_d/2.d0 + integer, optional :: errlen=512 + real(kind=8), parameter,optional :: gamma=0.0065d0 + end module constants + subroutine dcomputepi(pi,pressure,nx,ny,nz) ! in :_wrffortran:wrf_user.f90 + threadsafe + use constants, only: p1000mb,cp,r_d + real(kind=8) dimension(nx,ny,nz),intent(out,in) :: pi + real(kind=8) dimension(nx,ny,nz),intent(in),depend(nx,ny,nz) :: pressure + integer, optional,intent(in),check(shape(pi,0)==nx),depend(pi) :: nx=shape(pi,0) + integer, optional,intent(in),check(shape(pi,1)==ny),depend(pi) :: ny=shape(pi,1) + integer, optional,intent(in),check(shape(pi,2)==nz),depend(pi) :: nz=shape(pi,2) + end subroutine dcomputepi + subroutine dcomputetk(tk,pressure,theta,nx) ! in :_wrffortran:wrf_user.f90 + threadsafe + use constants, only: p1000mb,cp,r_d + real(kind=8) dimension(nx),intent(out,in) :: tk + real(kind=8) dimension(nx),intent(in),depend(nx) :: pressure + real(kind=8) dimension(nx),intent(in),depend(nx) :: theta + integer, optional,intent(in),check(len(tk)>=nx),depend(tk) :: nx=len(tk) + end subroutine dcomputetk + subroutine dinterp3dz(data3d,out2d,zdata,desiredloc,nx,ny,nz,missingval) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx,ny,nz),intent(in) :: data3d + real(kind=8) dimension(nx,ny),intent(out,in),depend(nx,ny) :: out2d + real(kind=8) dimension(nx,ny,nz),intent(in),depend(nx,ny,nz) :: zdata + real(kind=8) intent(in) :: desiredloc + integer, optional,intent(in),check(shape(data3d,0)==nx),depend(data3d) :: nx=shape(data3d,0) + integer, optional,intent(in),check(shape(data3d,1)==ny),depend(data3d) :: ny=shape(data3d,1) + integer, optional,intent(in),check(shape(data3d,2)==nz),depend(data3d) :: nz=shape(data3d,2) + real(kind=8) intent(in) :: missingval + end subroutine dinterp3dz + subroutine dzstag(znew,nx,ny,nz,z,nxz,nyz,nzz,terrain) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx,ny,nz),intent(out,in) :: znew + integer, optional,intent(in),check(shape(znew,0)==nx),depend(znew) :: nx=shape(znew,0) + integer, optional,intent(in),check(shape(znew,1)==ny),depend(znew) :: ny=shape(znew,1) + integer, optional,intent(in),check(shape(znew,2)==nz),depend(znew) :: nz=shape(znew,2) + real(kind=8) dimension(nxz,nyz,nzz),intent(in) :: z + integer, optional,intent(in),check(shape(z,0)==nxz),depend(z) :: nxz=shape(z,0) + integer, optional,intent(in),check(shape(z,1)==nyz),depend(z) :: nyz=shape(z,1) + integer, optional,intent(in),check(shape(z,2)==nzz),depend(z) :: nzz=shape(z,2) + real(kind=8) dimension(nxz,nyz),intent(in),depend(nxz,nyz) :: terrain + end subroutine dzstag + subroutine dinterp2dxy(v3d,v2d,xy,nx,ny,nz,nxy) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx,ny,nz),intent(in) :: v3d + real(kind=8) dimension(nxy,nz),intent(out,in),depend(nz) :: v2d + real(kind=8) dimension(2,nxy),intent(in),depend(nxy) :: xy + integer, optional,intent(in),check(shape(v3d,0)==nx),depend(v3d) :: nx=shape(v3d,0) + integer, optional,intent(in),check(shape(v3d,1)==ny),depend(v3d) :: ny=shape(v3d,1) + integer, optional,intent(in),check(shape(v3d,2)==nz),depend(v3d) :: nz=shape(v3d,2) + integer, optional,intent(in),check(shape(v2d,0)==nxy),depend(v2d) :: nxy=shape(v2d,0) + end subroutine dinterp2dxy + subroutine dinterp1d(v_in,v_out,z_in,z_out,vmsg,nz_in,nz_out) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nz_in),intent(in) :: v_in + real(kind=8) dimension(nz_out),intent(out,in) :: v_out + real(kind=8) dimension(nz_in),intent(in),depend(nz_in) :: z_in + real(kind=8) dimension(nz_out),intent(in),depend(nz_out) :: z_out + real(kind=8) intent(in) :: vmsg + integer, optional,intent(in),check(len(v_in)>=nz_in),depend(v_in) :: nz_in=len(v_in) + integer, optional,intent(in),check(len(v_out)>=nz_out),depend(v_out) :: nz_out=len(v_out) + end subroutine dinterp1d + subroutine dcomputeseaprs(nx,ny,nz,z,t,p,q,sea_level_pressure,t_sea_level,t_surf,level,errstat,errmsg) ! in :_wrffortran:wrf_user.f90 + threadsafe + use constants, only: r,errlen,gamma,g + integer, optional,intent(in),check(shape(z,0)==nx),depend(z) :: nx=shape(z,0) + integer, optional,intent(in),check(shape(z,1)==ny),depend(z) :: ny=shape(z,1) + integer, optional,intent(in),check(shape(z,2)==nz),depend(z) :: nz=shape(z,2) + real(kind=8) dimension(nx,ny,nz),intent(in) :: z + real(kind=8) dimension(nx,ny,nz),intent(in),depend(nx,ny,nz) :: t + real(kind=8) dimension(nx,ny,nz),intent(in),depend(nx,ny,nz) :: p + real(kind=8) dimension(nx,ny,nz),intent(in),depend(nx,ny,nz) :: q + real(kind=8) dimension(nx,ny),intent(out,in),depend(nx,ny) :: sea_level_pressure + real(kind=8) dimension(nx,ny),intent(inout),depend(nx,ny) :: t_sea_level + real(kind=8) dimension(nx,ny),intent(inout),depend(nx,ny) :: t_surf + integer dimension(nx,ny),intent(inout),depend(nx,ny) :: level + integer, optional,intent(inout) :: errstat + character*errlen, optional,intent(inout) :: errmsg + end subroutine dcomputeseaprs + subroutine dfilter2d(a,b,nx,ny,it,missing) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx,ny),intent(in,out) :: a + real(kind=8) dimension(nx,ny),intent(inout),depend(nx,ny) :: b + integer, optional,intent(in),check(shape(a,0)==nx),depend(a) :: nx=shape(a,0) + integer, optional,intent(in),check(shape(a,1)==ny),depend(a) :: ny=shape(a,1) + integer intent(in) :: it + real(kind=8) intent(in) :: missing + end subroutine dfilter2d + subroutine filter2d(a,b,nx,ny,it,missing) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=4) dimension(nx,ny),intent(in,out) :: a + real(kind=4) dimension(nx,ny),intent(inout),depend(nx,ny) :: b + integer, optional,intent(in),check(shape(a,0)==nx),depend(a) :: nx=shape(a,0) + integer, optional,intent(in),check(shape(a,1)==ny),depend(a) :: ny=shape(a,1) + integer intent(in) :: it + real(kind=4) intent(in) :: missing + end subroutine filter2d + subroutine dcomputerh(qv,p,t,rh,nx) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx),intent(in) :: qv + real(kind=8) dimension(nx),intent(in),depend(nx) :: p + real(kind=8) dimension(nx),intent(in),depend(nx) :: t + real(kind=8) dimension(nx),intent(out,in),depend(nx) :: rh + integer, optional,intent(in),check(len(qv)>=nx),depend(qv) :: nx=len(qv) + end subroutine dcomputerh + subroutine dgetijlatlong(lat_array,long_array,lat,longitude,ii,jj,nx,ny,imsg) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx,ny),intent(in) :: lat_array + real(kind=8) dimension(nx,ny),intent(in),depend(nx,ny) :: long_array + real(kind=8) :: lat + real(kind=8) :: longitude + integer intent(out,in) :: ii + integer intent(out,in) :: jj + integer, optional,intent(in),check(shape(lat_array,0)==nx),depend(lat_array) :: nx=shape(lat_array,0) + integer, optional,intent(in),check(shape(lat_array,1)==ny),depend(lat_array) :: ny=shape(lat_array,1) + integer intent(in) :: imsg + end subroutine dgetijlatlong + subroutine dcomputeuvmet(u,v,uvmet,longca,longcb,flong,flat,cen_long,cone,rpd,nx,ny,nz,nxp1,nyp1,istag,is_msg_val,umsg,vmsg,uvmetmsg) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nxp1,ny,nz),intent(in) :: u + real(kind=8) dimension(nx,nyp1,nz),intent(in),depend(nz) :: v + real(kind=8) dimension(nx,ny,nz,2),intent(out,in),depend(nx,ny,nz) :: uvmet + real(kind=8) dimension(nx,ny),intent(inout),depend(nx,ny) :: longca + real(kind=8) dimension(nx,ny),intent(inout),depend(nx,ny) :: longcb + real(kind=8) dimension(nx,ny),intent(in),depend(nx,ny) :: flong + real(kind=8) dimension(nx,ny),intent(in),depend(nx,ny) :: flat + real(kind=8) intent(in) :: cen_long + real(kind=8) intent(in) :: cone + real(kind=8) intent(in) :: rpd + integer, optional,intent(in),check(shape(v,0)==nx),depend(v) :: nx=shape(v,0) + integer, optional,intent(in),check(shape(u,1)==ny),depend(u) :: ny=shape(u,1) + integer, optional,intent(in),check(shape(u,2)==nz),depend(u) :: nz=shape(u,2) + integer, optional,intent(in),check(shape(u,0)==nxp1),depend(u) :: nxp1=shape(u,0) + integer, optional,intent(in),check(shape(v,1)==nyp1),depend(v) :: nyp1=shape(v,1) + integer intent(in) :: istag + logical intent(in) :: is_msg_val + real(kind=8) intent(in) :: umsg + real(kind=8) intent(in) :: vmsg + real(kind=8) intent(in) :: uvmetmsg + end subroutine dcomputeuvmet + subroutine dcomputetd(td,pressure,qv_in,nx) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx),intent(out,in) :: td + real(kind=8) dimension(nx),intent(in),depend(nx) :: pressure + real(kind=8) dimension(nx),intent(in),depend(nx) :: qv_in + integer, optional,intent(in),check(len(td)>=nx),depend(td) :: nx=len(td) + end subroutine dcomputetd + subroutine dcomputeiclw(iclw,pressure,qc_in,nx,ny,nz) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx,ny),intent(out,in) :: iclw + real(kind=8) dimension(nx,ny,nz),intent(in),depend(nx,ny) :: pressure + real(kind=8) dimension(nx,ny,nz),intent(in),depend(nx,ny,nz) :: qc_in + integer, optional,intent(in),check(shape(iclw,0)==nx),depend(iclw) :: nx=shape(iclw,0) + integer, optional,intent(in),check(shape(iclw,1)==ny),depend(iclw) :: ny=shape(iclw,1) + integer, optional,intent(in),check(shape(pressure,2)==nz),depend(pressure) :: nz=shape(pressure,2) + end subroutine dcomputeiclw + subroutine testfunc(a,b,nx,ny,nz,errstat,errstr) ! in :_wrffortran:wrf_user.f90 + threadsafe + real(kind=8) dimension(nx,ny,nz),intent(in) :: a + real(kind=8) dimension(nx,ny,nz),intent(out,in),depend(nx,ny,nz) :: b + integer, optional,intent(in),check(shape(a,0)==nx),depend(a) :: nx=shape(a,0) + integer, optional,intent(in),check(shape(a,1)==ny),depend(a) :: ny=shape(a,1) + integer, optional,intent(in),check(shape(a,2)==nz),depend(a) :: nz=shape(a,2) + integer, optional,intent(inout) :: errstat + character*512, optional,intent(inout) :: errstr + end subroutine testfunc + end interface +end python module _wrffortran + +! This file was auto-generated with f2py (version:2). +! See http://cens.ioc.ee/projects/f2py2e/ diff --git a/setup.py b/setup.py index 66f3ac0..16f257e 100755 --- a/setup.py +++ b/setup.py @@ -13,11 +13,18 @@ ext2 = numpy.distutils.core.Extension( "src/wrf/wrfcape.pyf"] ) +ext3 = numpy.distutils.core.Extension( + name = "wrf._wrffortran", + sources = ["fortran/constants.f90", + "fortran/wrf_user.f90", + "fortran/wrffortran.pyf"] + ) + numpy.distutils.core.setup( name = "wrf", version = "0.0.1", packages = setuptools.find_packages("src"), - ext_modules = [ext1,ext2], + ext_modules = [ext1,ext2,ext3], package_dir={"":"src"}, #namespace_packages=["wrf"], scripts=[], diff --git a/src/wrf/cape.py b/src/wrf/cape.py index 31e8292..1809b4d 100755 --- a/src/wrf/cape.py +++ b/src/wrf/cape.py @@ -4,7 +4,8 @@ from __future__ import (absolute_import, division, print_function, import numpy as n import numpy.ma as ma -from .extension import computetk,computecape +#from .extension import computetk,computecape +from .extension import _tk,computecape from .destag import destagger from .constants import Constants, ConversionFactors from .util import extract_vars, combine_with @@ -39,7 +40,7 @@ def get_2dcape(wrfnc, timeidx=0, method="cat", full_t = t + Constants.T_BASE full_p = p + pb - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) geopt = ph + phb geopt_unstag = destagger(geopt, -3) @@ -101,7 +102,7 @@ def get_3dcape(wrfnc, timeidx=0, method="cat", full_t = t + Constants.T_BASE full_p = p + pb - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) geopt = ph + phb geopt_unstag = destagger(geopt, -3) diff --git a/src/wrf/constants.py b/src/wrf/constants.py index 67bff0b..e096b42 100755 --- a/src/wrf/constants.py +++ b/src/wrf/constants.py @@ -14,6 +14,7 @@ class Constants(object): PI = 3.141592653589793 DEFAULT_FILL = 9.9692099683868690e+36 # Value is from netcdf.h WRF_EARTH_RADIUS = 6370000. + ERRLEN = 512 class ConversionFactors(object): PA_TO_HPA = .01 diff --git a/src/wrf/ctt.py b/src/wrf/ctt.py index 7fb5534..58db5f9 100644 --- a/src/wrf/ctt.py +++ b/src/wrf/ctt.py @@ -3,7 +3,8 @@ from __future__ import (absolute_import, division, print_function, import numpy as n -from .extension import computectt, computetk +#from .extension import computectt, computetk +from .extension import computectt, _tk from .constants import Constants, ConversionFactors from .destag import destagger from .decorators import convert_units @@ -55,7 +56,7 @@ def get_ctt(wrfnc, timeidx=0, method="cat", full_p = p + pb p_hpa = full_p * ConversionFactors.PA_TO_HPA full_t = t + Constants.T_BASE - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) geopt = ph + phb geopt_unstag = destagger(geopt, -3) diff --git a/src/wrf/dbz.py b/src/wrf/dbz.py index 4cd4aa2..a08603f 100755 --- a/src/wrf/dbz.py +++ b/src/wrf/dbz.py @@ -3,7 +3,8 @@ from __future__ import (absolute_import, division, print_function, import numpy as n -from .extension import computedbz,computetk +#from .extension import computedbz,computetk +from .extension import computedbz, _tk from .constants import Constants from .util import extract_vars from .metadecorators import copy_and_set_metadata @@ -58,7 +59,7 @@ def get_dbz(wrfnc, timeidx=0, method="cat", full_t = t + Constants.T_BASE full_p = p + pb - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) ivarint = 0 if do_varint: diff --git a/src/wrf/decorators.py b/src/wrf/decorators.py index 443c09f..c2eaf7c 100644 --- a/src/wrf/decorators.py +++ b/src/wrf/decorators.py @@ -1,14 +1,15 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -from collections import Iterable +from collections import Iterable, OrderedDict import wrapt import numpy as np import numpy.ma as ma from .units import do_conversion, check_units -from .util import (iter_left_indexes, viewitems, from_args, npvalues, py3range) +from .util import (iter_left_indexes, viewitems, viewvalues, from_args, + npvalues, py3range, combine_dims, isstr) from .config import xarray_enabled if xarray_enabled(): @@ -86,7 +87,7 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, return wrapped(*args, **kwargs) # Start by getting the left-most 'extra' dims - extra_dims = [ref_var_shape[x] for x in py3range(extra_dim_num)] + extra_dims = ref_var_shape[0:extra_dim_num] out_inited = False for left_idxs in iter_left_indexes(extra_dims): @@ -106,48 +107,48 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, for key,val in viewitems(kwargs)} # Call the numerical routine - res = wrapped(*new_args, **new_kargs) + result = wrapped(*new_args, **new_kargs) - if isinstance(res, np.ndarray): + if isinstance(result, np.ndarray): # Output array if not out_inited: - outdims = _calc_out_dims(res, extra_dims) - if not isinstance(res, ma.MaskedArray): + outdims = _calc_out_dims(result, extra_dims) + if not isinstance(result, ma.MaskedArray): output = np.empty(outdims, ref_var.dtype) masked = False else: output = ma.MaskedArray( np.zeros(outdims, ref_var.dtype), mask=np.zeros(outdims, np.bool_), - fill_value=res.fill_value) + fill_value=result.fill_value) masked = True out_inited = True if not masked: - output[left_and_slice_idxs] = res[:] + output[left_and_slice_idxs] = result[:] else: - output.data[left_and_slice_idxs] = res.data[:] - output.mask[left_and_slice_idxs] = res.mask[:] + output.data[left_and_slice_idxs] = result.data[:] + output.mask[left_and_slice_idxs] = result.mask[:] else: # This should be a list or a tuple (cape) if not out_inited: - outdims = _calc_out_dims(res[0], extra_dims) - if not isinstance(res[0], ma.MaskedArray): + outdims = _calc_out_dims(result[0], extra_dims) + if not isinstance(result[0], ma.MaskedArray): output = [np.empty(outdims, ref_var.dtype) - for i in py3range(len(res))] + for i in py3range(len(result))] masked = False else: output = [ma.MaskedArray( np.zeros(outdims, ref_var.dtype), mask=np.zeros(outdims, np.bool_), - fill_value=res[0].fill_value) - for i in py3range(len(res))] + fill_value=result[0].fill_value) + for i in py3range(len(result))] masked = True out_inited = True - for i,outarr in enumerate(res): + for i,outarr in enumerate(result): if not masked: output[i][left_and_slice_idxs] = outarr[:] else: @@ -158,9 +159,168 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, return output return func_wrapper + + + +def left_iter_nocopy(ref_var_expected_dims, + ref_var_right_ndims, + insert_dims=None, + ref_var_idx=None, + ref_var_name=None, + ignore_args=None, + ignore_kargs=None, + outviews="outview", + alg_dtype=np.float64, + cast_output=True): + """Decorator to handle iterating over leftmost dimensions when using + multiple files and/or multiple times. + + Arguments: + - ref_var_expected_dims - the number of dimensions that the Fortran + algorithm is expecting for the reference variable + - ref_var_right_ndims - the number of ndims from the right to keep for + the reference variable when making the output. Can also be a + combine_dims instance if the sizes are determined from multiple + variables. + - insert_dims - a sequence of dimensions to insert between the left + dimenions (e.g. time) and the kept right dimensions + - ref_var_idx - the index in args used as the reference variable for + calculating leftmost dimensions + - ref_var_name - the keyword argument name for kwargs used as the + reference varible for calculating leftmost dimensions + - alg_out_fixed_dims - additional fixed dimension sizes for the + numerical algorithm (e.g. uvmet has a fixed left dimsize of 2) + - ignore_args - indexes of any arguments which should be passed + directly without any slicing + - ignore_kargs - keys of any keyword arguments which should be passed + directly without slicing + - outviews - a single key or sequence of keys indicating the keyword + argument used as the output variable(s) + - algtype - the data type used in the numerical routine + - cast_output - cast the final output to the ref_var data type + + + """ + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + _ignore_args = ignore_args if ignore_args is not None else () + _ignore_kargs = ignore_kargs if ignore_kargs is not None else () + _outkeys = [outviews] if isstr(outviews) else outviews + + if ref_var_idx is not None: + ref_var = args[ref_var_idx] + else: + ref_var = kwargs[ref_var_name] + + ref_var_dtype = ref_var.dtype + ref_var_shape = ref_var.shape + extra_dim_num = ref_var.ndim - ref_var_expected_dims + + # No special left side iteration, return the function result + if (extra_dim_num == 0): + return wrapped(*args, **kwargs) + + # Start by getting the left-most 'extra' dims + extra_dims = ref_var_shape[0:extra_dim_num] + + mid_dims = () if insert_dims is None else tuple(insert_dims) + + if not isinstance(ref_var_right_ndims, combine_dims): + right_dims = ref_var_shape[-ref_var_right_ndims:] + else: + right_dims = ref_var_right_ndims(*args) + + left_dims = extra_dims + + outdims = left_dims + mid_dims + right_dims + + if "outview" not in kwargs: + outd = OrderedDict((outkey, np.empty(outdims, alg_dtype)) + for outkey in _outkeys) + + for left_idxs in iter_left_indexes(extra_dims): + # Make the left indexes plus a single slice object + # The single slice will handle all the dimensions to + # the right (e.g. [1,1,:]) + left_and_slice_idxs = left_idxs + (slice(None), ) + + + # Slice the args if applicable + new_args = [arg[left_and_slice_idxs] + if i not in _ignore_args else arg + for i,arg in enumerate(args)] + + # Slice the kwargs if applicable + new_kargs = {key:(val[left_and_slice_idxs] + if key not in _ignore_kargs else val) + for key,val in viewitems(kwargs)} + + + # Insert the output views if one hasn't been provided + if "outview" not in new_kargs: + for outkey,output in viewitems(outd): + outview = output[left_and_slice_idxs] + new_kargs[outkey] = outview + + result = wrapped(*new_args, **new_kargs) + + # Make sure the result is the same data as what got passed in + # Can delete this once everything works + if (result.__array_interface__["data"][0] != + outview.__array_interface__["data"][0]): + raise RuntimeError("output array was copied") + + if len(outd) == 1: + output = next(iter(viewvalues(outd))) + else: + output = tuple(arr for arr in viewvalues(outd)) + + if cast_output: + if isinstance(output, np.ndarray): + output = output.astype(ref_var_dtype) + else: + output = tuple(arr.astype(ref_var_dtype) for arr in output) + + return output + + return func_wrapper +# Only flip the input arguments, but flip the result +def flip_vertical(argidxs=None, argkeys=None, flip_result=True): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + + _argidxs = argidxs if argidxs is not None else () + _argkeys = argkeys if argkeys is not None else () + + # Slice the args if applicable + new_args = [np.ascontiguousarray(arg[...,::-1,:,:]) + if i in _argidxs else arg + for i,arg in enumerate(args)] + + # Slice the kwargs if applicable + new_kargs = {key:np.ascontiguousarray(kwarg[...,::-1,:,:]) + if key in _argkeys else kwarg + for key,kwarg in viewitems(kwargs)} + + + result = wrapped(*args, **kwargs) + + if flip_result: + if isinstance(result, np.ndarray): + return np.ascontiguousarray(result[...,::-1,:,:]) + else: + # Sequence + return [np.ascontiguousarray(arr[...,::-1,:,:]) + for arr in result] + + return result + + return func_wrapper -def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=np.float64): +def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, + alg_dtype=np.float64, + outviews="outview"): """Decorator to handle casting to/from required dtype used in numerical routine. @@ -170,31 +330,47 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=np.float64): _arg_idxs = arg_idxs if arg_idxs is not None else () _karg_names = karg_names if karg_names is not None else () + # Handle output views if applicable + _outkeys = [outviews] if isstr(outviews) else _outviews + _outviews = from_args(wrapped, _outkeys, *args, **kwargs) + + has_outview = False + for outkey in _outkeys: + _outview = _outviews[outkey] + if _outview is not None: + has_outview = True + + orig_type = args[ref_idx].dtype - new_args = [arg.astype(dtype) + new_args = [arg.astype(alg_dtype) if i in _arg_idxs else arg for i,arg in enumerate(args)] - new_kargs = {key:(val.astype(dtype) + new_kargs = {key:(val.astype(alg_dtype) if key in _karg_names else val) for key,val in viewitems(kwargs)} - res = wrapped(*new_args, **new_kargs) + result = wrapped(*new_args, **new_kargs) - if isinstance(res, np.ndarray): - if res.dtype == orig_type: - return res - return res.astype(orig_type) - else: # got back a sequence of arrays - return tuple(arr.astype(orig_type) - if arr.dtype != orig_type else arr - for arr in res) + # Do nothing for supplied output views + if not has_outview: + if isinstance(result, np.ndarray): + if result.dtype == orig_type: + return result + return result.astype(orig_type) + else: # got back a sequence of arrays + return tuple(arr.astype(orig_type) + if arr.dtype != orig_type else arr + for arr in result) + + return result return func_wrapper def _extract_and_transpose(arg, do_transpose): + if xarray_enabled(): if isinstance(arg, DataArray): arg = npvalues(arg) @@ -207,7 +383,7 @@ def _extract_and_transpose(arg, do_transpose): return arg -def handle_extract_transpose(do_transpose=True): +def handle_extract_transpose(do_transpose=True, outviews="outview"): """Decorator to extract the data array from a DataArray and also transposes the view of the data if the data is not fortran contiguous. @@ -215,21 +391,35 @@ def handle_extract_transpose(do_transpose=True): @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): + # Handle output views if applicable + _outkeys = [outviews] if isstr(outviews) else _outviews + _outviews = from_args(wrapped, _outkeys, *args, **kwargs) + + has_outview = False + for outkey in _outkeys: + _outview = _outviews[outkey] + if _outview is not None: + has_outview = True + new_args = [_extract_and_transpose(arg, do_transpose) for arg in args] - new_kargs = {key:_extract_and_transpose(val) + new_kargs = {key:_extract_and_transpose(val, do_transpose) for key,val in viewitems(kwargs)} - res = wrapped(*new_args, **new_kargs) + result = wrapped(*new_args, **new_kargs) + + # Do nothing for supplied output views + if has_outview: + return result - if isinstance(res, np.ndarray): - if res.flags.f_contiguous and res.ndim > 1: - return res.T - elif isinstance(res, Iterable): + if isinstance(result, np.ndarray): + if result.flags.f_contiguous and result.ndim > 1: + return result.T + elif isinstance(result, Iterable): return tuple(x.T if x.flags.f_contiguous and x.ndim > 1 else x - for x in res) + for x in result) - return res + return result return func_wrapper diff --git a/src/wrf/dewpoint.py b/src/wrf/dewpoint.py index cc6c317..e2659b8 100755 --- a/src/wrf/dewpoint.py +++ b/src/wrf/dewpoint.py @@ -1,7 +1,8 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -from .extension import computetd +#from .extension import computetd +from .extension import _td from .decorators import convert_units from .metadecorators import copy_and_set_metadata from .util import extract_vars @@ -26,7 +27,7 @@ def get_dp(wrfnc, timeidx=0, method="cat", squeeze=True, full_p = .01*(p + pb) qvapor[qvapor < 0] = 0 - td = computetd(full_p, qvapor) + td = _td(full_p, qvapor) return td @copy_and_set_metadata(copy_varname="Q2", name="td2", @@ -44,7 +45,7 @@ def get_dp_2m(wrfnc, timeidx=0, method="cat", squeeze=True, q2 = ncvars["Q2"] q2[q2 < 0] = 0 - td = computetd(psfc, q2) + td = _td(psfc, q2) return td diff --git a/src/wrf/extension.py b/src/wrf/extension.py index 20ff98e..034190f 100755 --- a/src/wrf/extension.py +++ b/src/wrf/extension.py @@ -5,6 +5,7 @@ import numpy as np from .constants import Constants from .psadlookup import get_lookup_tables +# Old way from ._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, f_computeslp, f_computetk, f_computetd, f_computerh, f_computeabsvort,f_computepvo, f_computeeth, @@ -14,136 +15,307 @@ from ._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, f_lltoij, f_ijtoll, f_converteta, f_computectt, f_monotonic, f_filter2d, f_vintrp) from ._wrfcape import f_computecape + +# New way +from ._wrffortran import (dcomputetk, dinterp3dz, dinterp2dxy, dinterp1d, + dcomputeseaprs, dfilter2d, dcomputerh, dcomputeuvmet, + dcomputetd) + from .decorators import (handle_left_iter, handle_casting, handle_extract_transpose) -from .util import py3range -from .uvdecorator import uvmet_left_iter - -__all__ = ["FortranException", "computeslp", "computetk", "computetd", - "computerh", "computeavo", "computepvo", "computeeth", - "computeuvmet","computeomega", "computetv", "computesrh", - "computeuh", "computepw","computedbz","computecape", - "computeij", "computell", "computeeta", "computectt", - "interp2dxy", "interpz3d", "interp1d", "computeinterpline", - "computevertcross"] +from .decorators import (left_iter_nocopy) +from .util import py3range, combine_dims, _npbytes_to_str +from .uvdecorator import uvmet_left_iter, uvmet_left_iter_nocopy + +__all__ = [] +# __all__ += ["FortranException", "computeslp", "computetk", "computetd", +# "computerh", "computeavo", "computepvo", "computeeth", +# "computeuvmet","computeomega", "computetv", "computesrh", +# "computeuh", "computepw","computedbz","computecape", +# "computeij", "computell", "computeeta", "computectt", +# "interp2dxy", "interpz3d", "interp1d", "computeinterpline", +# "computevertcross"] +# __all__ += ["", ""] class FortranException(Exception): def __call__(self, message): raise self.__class__(message) -@handle_left_iter(3,0, ignore_args=(2,3)) +# IMPORTANT! Unless otherwise noted, all variables used in the routines +# below assume that fortran-ordering views are being used. + + + +# @handle_left_iter(3,0, ignore_args=(2,3)) +# @handle_casting(arg_idxs=(0,1)) +# @handle_extract_transpose() +# def interpz3d(field3d, z, desiredloc, missingval): +# result = f_interpz3d(field3d, +# z, +# desiredloc, +# missingval) +# return result + +@left_iter_nocopy(3, 2, ref_var_idx=0, ignore_args=(2,3)) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def interpz3d(field3d, z, desiredloc, missingval): - res = f_interpz3d(field3d, - z, - desiredloc, - missingval) - return res - -@handle_left_iter(3,0, ignore_args=(1,)) +def _interpz3d(field3d, z, desiredloc, missingval, outview=None): + if outview is None: + outview = np.empty(field3d.shape[0:2], np.float64, order="F") + + result = dinterp3dz(field3d, + outview, + z, + desiredloc, + missingval) + return result + +# @handle_left_iter(3,0, ignore_args=(1,)) +# @handle_casting(arg_idxs=(0,1)) +# @handle_extract_transpose() +# def interp2dxy(field3d, xy): +# result = f_interp2dxy(field3d, +# xy) +# return result + +@left_iter_nocopy(3, combine_dims([(0,-3),(1,-2)]), ref_var_idx=0, + ignore_args=(1,)) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def interp2dxy(field3d, xy): - res = f_interp2dxy(field3d, - xy) - return res - -@handle_left_iter(1, 0, ignore_args=(2,3)) +def _interp2dxy(field3d, xy, outview=None): + if outview is None: + outview = np.empty((xy.shape[-1], field3d.shape[-1]), np.float64, + order="F") + + result = dinterp2dxy(field3d, + outview, + xy) + return result + +# @handle_left_iter(1, 0, ignore_args=(2,3)) +# @handle_casting(arg_idxs=(0,1,2)) +# @handle_extract_transpose() +# def interp1d(v_in, z_in, z_out, missingval): +# result = f_interp1d(v_in, +# z_in, +# z_out, +# missingval) +# +# return result + +@left_iter_nocopy(1, 1, ref_var_idx=0, ignore_args=(2,3)) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() -def interp1d(v_in, z_in, z_out, missingval): - res = f_interp1d(v_in, - z_in, - z_out, - missingval) - - return res - -@handle_left_iter(3, 0, ignore_args=(1,4,3)) +def _interp1d(v_in, z_in, z_out, missingval, outview=None): + + if outview is None: + outview = np.empty_like(z_out) + + result = dinterp1d(v_in, + outview, + z_in, + z_out, + missingval) + + return result + +# @handle_left_iter(3, 0, ignore_args=(1,4,3)) +# @handle_casting(arg_idxs=(0,)) +# @handle_extract_transpose(do_transpose=False) +# def computevertcross(field3d, xy, var2dz, z_var2d, missingval): +# var2d = np.empty((z_var2d.size, xy.shape[0]), dtype=var2dz.dtype) +# var2dtmp = interp2dxy(field3d, xy) +# +# for i in py3range(xy.shape[0]): +# var2d[:,i] = interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, missingval) +# +# return var2d + +@left_iter_nocopy(3, combine_dims([(3,0), (1,0)]), + ref_var_idx=0, ignore_args=(1,3,4)) @handle_casting(arg_idxs=(0,)) @handle_extract_transpose(do_transpose=False) -def computevertcross(field3d, xy, var2dz, z_var2d, missingval): - var2d = np.empty((z_var2d.size, xy.shape[0]), dtype=var2dz.dtype) - var2dtmp = interp2dxy(field3d, xy) +def _vertcross(field3d, xy, var2dz, z_var2d, missingval, outview=None): + # Note: This is using C-ordering + if outview is None: + outview = np.empty((z_var2d.shape[0], xy.shape[0]), dtype=var2dz.dtype) + + var2dtmp = _interp2dxy(field3d, xy) for i in py3range(xy.shape[0]): - var2d[:,i] = interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, missingval) - - return var2d - -@handle_left_iter(2, 0, ignore_args=(1,)) + outview[:,i] = _interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, + missingval) + + return outview + +# @handle_left_iter(2, 0, ignore_args=(1,)) +# @handle_casting(arg_idxs=(0,)) +# @handle_extract_transpose(do_transpose=False) +# def interpline(field2d, xy): +# +# tmp_shape = [1] + [x for x in field2d.shape] +# var2dtmp = np.empty(tmp_shape, field2d.dtype) +# var2dtmp[0,:,:] = field2d[:,:] +# +# var1dtmp = interp2dxy(var2dtmp, xy) +# +# return var1dtmp[0,:] + + +@left_iter_nocopy(2, combine_dims([(1,0)]), ref_var_idx=0, ignore_args=(1,)) @handle_casting(arg_idxs=(0,)) @handle_extract_transpose(do_transpose=False) -def computeinterpline(field2d, xy): +def _interpline(field2d, xy, outview=None): + # Note: This is using C-ordering + if outview is None: + outview = np.empty(xy.shape[0], dtype=field2d.dtype) - tmp_shape = [1] + [x for x in field2d.shape] + tmp_shape = (1,) + field2d.shape var2dtmp = np.empty(tmp_shape, field2d.dtype) var2dtmp[0,:,:] = field2d[:,:] - var1dtmp = interp2dxy(var2dtmp, xy) - - return var1dtmp[0,:] - -@handle_left_iter(3,0) + var1dtmp = _interp2dxy(var2dtmp, xy) + + outview[:] = var1dtmp[0, :] + + return outview + +# @handle_left_iter(3,0) +# @handle_casting(arg_idxs=(0,1,2,3)) +# @handle_extract_transpose() +# def computeslp(z, t, p, q): +# +# t_surf = np.zeros(z.shape[0:2], np.float64, order="F") +# t_sea_level = np.zeros(z.shape[0:2], np.float64, order="F") +# level = np.zeros(z.shape[0:2], np.int32, order="F") +# +# result = f_computeslp(z, +# t, +# p, +# q, +# t_sea_level, +# t_surf, +# level, +# FortranException()) +# +# return result + + +@left_iter_nocopy(3, 2, ref_var_idx=0) @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() -def computeslp(z, t, p, q): +def _slp(z, t, p, q, outview=None): t_surf = np.zeros(z.shape[0:2], np.float64, order="F") t_sea_level = np.zeros(z.shape[0:2], np.float64, order="F") level = np.zeros(z.shape[0:2], np.int32, order="F") - res = f_computeslp(z, - t, - p, - q, - t_sea_level, - t_surf, - level, - FortranException()) - - return res + if outview is None: + outview = np.empty(z.shape[0:2], np.float64, order="F") + + errstat = np.array(0) + errmsg = np.zeros(Constants.ERRLEN, "c") + + result = dcomputeseaprs(z, + t, + p, + q, + outview, + t_sea_level, + t_surf, + level, + errstat=errstat, + errmsg=errmsg) + + if int(errstat) != 0: + raise RuntimeError("".join(_npbytes_to_str(errmsg)).strip()) + + return result + +# @handle_left_iter(3,0) +# @handle_casting(arg_idxs=(0,1)) +# @handle_extract_transpose() +# def computetk(pres, theta): +# # No need to transpose here since operations on 1D array +# shape = pres.shape +# result = f_computetk(pres.ravel(order="A"), +# theta.ravel(order="A")) +# result = np.reshape(result, shape, order="F") +# +# return result -@handle_left_iter(3,0) + +# Note: No left iteration decorator needed with 1D arrays @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def computetk(pres, theta): +def _tk(pressure, theta, outview=None): # No need to transpose here since operations on 1D array - shape = pres.shape - res = f_computetk(pres.ravel(order="A"), - theta.ravel(order="A")) - res = np.reshape(res, shape, order="F") + shape = pressure.shape + if outview is None: + outview = np.empty_like(pressure) + result = dcomputetk(outview.ravel(order="A"), + pressure.ravel(order="A"), + theta.ravel(order="A")) + result = np.reshape(result, shape, order="F") - return res + return result + +# Note: No left iteration decorator needed with 1D arrays +# @handle_casting(arg_idxs=(0,1)) +# @handle_extract_transpose() +# def computetd(pressure, qv_in): +# shape = pressure.shape +# result = f_computetd(pressure.ravel(order="A"), +# qv_in.ravel(order="A")) +# result = np.reshape(result, shape, order="F") +# return result # Note: No left iteration decorator needed with 1D arrays @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def computetd(pressure, qv_in): +def _td(pressure, qv_in, outview=None): shape = pressure.shape - res = f_computetd(pressure.ravel(order="A"), - qv_in.ravel(order="A")) - res = np.reshape(res, shape, order="F") - return res - -# Note: No decorator needed with 1D arrays + if outview is None: + outview = np.empty_like(pressure) + result = dcomputetd(outview.ravel(order="A"), + pressure.ravel(order="A"), + qv_in.ravel(order="A")) + result = np.reshape(result, shape, order="F") + + return result + +# # Note: No decorator needed with 1D arrays +# @handle_casting(arg_idxs=(0,1,2)) +# @handle_extract_transpose() +# def computerh(qv, q, t): +# shape = qv.shape +# result = f_computerh(qv.ravel(order="A"), +# q.ravel(order="A"), +# t.ravel(order="A")) +# result = np.reshape(result, shape, order="F") +# return result + +# Note: No left iteration decorator needed with 1D arrays @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() -def computerh(qv, q, t): +def _rh(qv, q, t, outview=None): shape = qv.shape - res = f_computerh(qv.ravel(order="A"), + if outview is None: + outview = np.empty_like(qv) + result = dcomputerh(qv.ravel(order="A"), q.ravel(order="A"), - t.ravel(order="A")) - res = np.reshape(res, shape, order="F") - return res + t.ravel(order="A"), + outview.ravel(order="A")) + result = np.reshape(result, shape, order="F") + + return result @handle_left_iter(3,0, ignore_args=(6,7)) @handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_extract_transpose() def computeavo(u, v, msfu, msfv, msfm, cor, dx, dy): - res = f_computeabsvort(u, + result = f_computeabsvort(u, v, msfu, msfv, @@ -152,14 +324,14 @@ def computeavo(u, v, msfu, msfv, msfm, cor, dx, dy): dx, dy) - return res + return result @handle_left_iter(3,2, ignore_args=(8,9)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6,7)) @handle_extract_transpose() def computepvo(u, v, theta, prs, msfu, msfv, msfm, cor, dx, dy): - res = f_computepvo(u, + result = f_computepvo(u, v, theta, prs, @@ -170,73 +342,111 @@ def computepvo(u, v, theta, prs, msfu, msfv, msfm, cor, dx, dy): dx, dy) - return res + return result @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() def computeeth(qv, tk, p): - res = f_computeeth(qv, + result = f_computeeth(qv, tk, p) - return res - -@uvmet_left_iter() + return result + +# @uvmet_left_iter() +# @handle_casting(arg_idxs=(0,1,2,3)) +# @handle_extract_transpose() +# def computeuvmet(u, v, lat, lon, cen_long, cone): +# longca = np.zeros((lat.shape[0], lat.shape[1]), np.float64, order="F") +# longcb = np.zeros((lon.shape[0], lon.shape[1]), np.float64, order="F") +# rpd = Constants.PI/180. +# +# +# # Make the 2D array a 3D array with 1 dimension +# if u.ndim < 3: +# u = u[:, :, np.newaxis] +# v = v[:, :, np.newaxis] +# +# # istag will always be false since winds are destaggered already +# # Missing values don't appear to be used, so setting to 0 +# result = f_computeuvmet(u, +# v, +# longca, +# longcb, +# lon, +# lat, +# cen_long, +# cone, +# rpd, +# 0, +# False, +# 0, +# 0, +# 0) +# +# +# return np.squeeze(result) + +# uvmet_left_iter needs to determine if the variable has missing values +@uvmet_left_iter_nocopy() @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() -def computeuvmet(u, v, lat, lon, cen_long, cone): - longca = np.zeros((lat.shape[0], lat.shape[1]), np.float64, order="F") - longcb = np.zeros((lon.shape[0], lon.shape[1]), np.float64, order="F") +def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False, + umissing=Constants.DEFAULT_FILL, vmissing=Constants.DEFAULT_FILL, + uvmetmissing=Constants.DEFAULT_FILL, outview=None): + longca = np.zeros(lat.shape[0:2], np.float64, order="F") + longcb = np.zeros(lon.shape[0:2], np.float64, order="F") rpd = Constants.PI/180. - # Make the 2D array a 3D array with 1 dimension if u.ndim < 3: - u = u.reshape((u.shape[0], u.shape[1], 1), order="F") - v = v.reshape((v.shape[0], v.shape[1], 1), order="F") - - # istag will always be false since winds are destaggered already - # Missing values don't appear to be used, so setting to 0 - res = f_computeuvmet(u, - v, - longca, - longcb, - lon, - lat, - cen_long, - cone, - rpd, - 0, - False, - 0, - 0, - 0) - - - return np.squeeze(res) + u = u[:, :, np.newaxis] + v = v[:, :, np.newaxis] + + if outview is None: + outdims = u.shape + (2,) + outview = np.empty(outdims, np.float64, order="F") + + result = dcomputeuvmet(u, + v, + outview, + longca, + longcb, + lon, + lat, + cen_long, + cone, + rpd, + isstag, + has_missing, + umissing, + vmissing, + uvmetmissing) + + return np.squeeze(result) @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() def computeomega(qv, tk, w, p): - res = f_computeomega(qv, + result = f_computeomega(qv, tk, w, p) - return res + return result @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() def computetv(tk, qv): - res = f_computetv(tk, + result = f_computetv(tk, qv) - return res + return result @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2)) @@ -245,7 +455,7 @@ def computewetbulb(p, tk, qv): PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITMK = PSADITMK.T - res = f_computewetbulb(p, + result = f_computewetbulb(p, tk, qv, PSADITHTE, @@ -253,20 +463,20 @@ def computewetbulb(p, tk, qv): PSADITMK, FortranException()) - return res + return result @handle_left_iter(3,0, ignore_args=(4,)) @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() def computesrh(u, v, z, ter, top): - res = f_computesrh(u, + result = f_computesrh(u, v, z, ter, top) - return res + return result @handle_left_iter(3,2, ignore_args=(5,6,7,8)) @handle_casting(arg_idxs=(0,1,2,3,4)) @@ -278,7 +488,7 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): tem2 = np.zeros((u.shape[0], u.shape[1], u.shape[2]), np.float64, order="F") - res = f_computeuh(zp, + result = f_computeuh(zp, mapfct, dx, dy, @@ -290,7 +500,7 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): tem1, tem2) - return res + return result @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2,3)) @@ -298,20 +508,20 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): def computepw(p, tv, qv, ht): zdiff = np.zeros((p.shape[0], p.shape[1]), np.float64, order="F") - res = f_computepw(p, + result = f_computepw(p, tv, qv, ht, zdiff) - return res + return result @handle_left_iter(3,0, ignore_args=(6,7,8)) @handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_extract_transpose() def computedbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin): - res = f_computedbz(p, + result = f_computedbz(p, tk, qv, qr, @@ -321,7 +531,13 @@ def computedbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin): ivarint, iliqskin) - return res + return result + + +# TODO: Make a new decorator to handle the flipping of the vertical +# The output arrays can be ignored and passed directly without flipping +# Then flip the final result. The copy is unavoidable, but at least it is +# only happening once, not every time. @handle_left_iter(3,0,ignore_args=(6,7,8)) @handle_casting(arg_idxs=(0,1,2,3,4,5)) @@ -338,7 +554,7 @@ def computecape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow): # along with tk,qv,and ht. # The extra mumbo-jumbo is so that the view created by numpy is fortran # contiguous. 'ascontiguousarray' only works in C ordering, hence the - # extra transposes. + # extra transposes. Note, this is probably making a copy flip_p = np.ascontiguousarray(p_hpa[:,:,::-1].T).T flip_tk = np.ascontiguousarray(tk[:,:,::-1].T).T flip_qv = np.ascontiguousarray(qv[:,:,::-1].T).T @@ -370,7 +586,7 @@ def computeij(map_proj, truelat1, truelat2, stdlon, lat1, lon1, pole_lat, pole_lon, knowni, knownj, dx, latinc, loninc, lat, lon): - res = f_lltoij(map_proj, + result = f_lltoij(map_proj, truelat1, truelat2, stdlon, @@ -387,13 +603,13 @@ def computeij(map_proj, truelat1, truelat2, stdlon, lon, FortranException()) - return res + return result def computell(map_proj, truelat1, truelat2, stdlon, lat1, lon1, pole_lat, pole_lon, knowni, knownj, dx, latinc, loninc, i, j): - res = f_ijtoll(map_proj, + result = f_ijtoll(map_proj, truelat1, truelat2, stdlon, @@ -410,7 +626,7 @@ def computell(map_proj, truelat1, truelat2, stdlon, lat1, lon1, j, FortranException()) - return res + return result @handle_left_iter(3,0, ignore_args=(3,)) @handle_casting(arg_idxs=(0,1,2)) @@ -420,7 +636,7 @@ def computeeta(full_t, znu, psfc, ptop): mean_t = np.zeros(full_t.shape, np.float64, order="F") temp_t = np.zeros(full_t.shape, np.float64, order="F") - res = f_converteta(full_t, + result = f_converteta(full_t, znu, psfc, ptop, @@ -428,13 +644,13 @@ def computeeta(full_t, znu, psfc, ptop): mean_t, temp_t) - return res + return result @handle_left_iter(3,0,ignore_args=(7,)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6)) @handle_extract_transpose() def computectt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci): - res = f_computectt(p_hpa, + result = f_computectt(p_hpa, tk, qice, qcld, @@ -443,45 +659,73 @@ def computectt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci): ter, haveqci) - return res - -@handle_left_iter(2,0,ignore_args=(1,)) + return result + +# @handle_left_iter(2,0,ignore_args=(1,)) +# @handle_casting(arg_idxs=(0,)) +# @handle_extract_transpose() +# def smooth2d(field, passes): +# # Unlike NCL, this routine will not modify the values in place, but +# # copies the original data before modifying it. This allows the decorator +# # to work properly and also behaves like the other methods. +# +# if isinstance(field, np.ma.MaskedArray): +# missing = field.fill_value +# else: +# missing = Constants.DEFAULT_FILL +# +# field_copy = field.copy(order="A") +# field_tmp = np.zeros(field_copy.shape, field_copy.dtype, order="F") +# +# f_filter2d(field_copy, +# field_tmp, +# missing, +# passes) +# +# # Don't transpose here since the fortran routine is not returning an +# # array. It's only modifying the existing array. +# return field_copy + +@left_iter_nocopy(2, 2, ref_var_idx=0, ignore_args=(1,)) @handle_casting(arg_idxs=(0,)) @handle_extract_transpose() -def smooth2d(field, passes): +def _smooth2d(field, passes, outview=None): # Unlike NCL, this routine will not modify the values in place, but - # copies the original data before modifying it. This allows the decorator - # to work properly and also behaves like the other methods. + # copies the original data before modifying it. if isinstance(field, np.ma.MaskedArray): missing = field.fill_value else: missing = Constants.DEFAULT_FILL - field_copy = field.copy(order="A") - field_tmp = np.zeros(field_copy.shape, field_copy.dtype, order="F") + if outview is None: + outview = field.copy(order="A") + else: + outview[:] = field[:] + + field_tmp = np.zeros(outview.shape, outview.dtype, order="F") - f_filter2d(field_copy, - field_tmp, - missing, - passes) + dfilter2d(outview, + field_tmp, + missing, + passes) # Don't transpose here since the fortran routine is not returning an # array. It's only modifying the existing array. - return field_copy + return outview @handle_left_iter(3,0,ignore_args=(3,4,5)) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() def monotonic(var, lvprs, coriolis, idir, delta, icorsw): - res = f_monotonic(var, + result = f_monotonic(var, lvprs, coriolis, idir, delta, icorsw) - return res + return result @handle_left_iter(3,0,ignore_args=(9,10,11,12,13,14)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6,7,8,9)) @@ -490,7 +734,7 @@ def vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp, vcarray, interp_levels, icase, extrap, vcor, logp, missing): - res = f_vintrp(field, + result = f_vintrp(field, pres, tk, qvp, @@ -506,7 +750,7 @@ def vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp, logp, missing) - return res + return result diff --git a/src/wrf/interp.py b/src/wrf/interp.py index 276bfce..e19e07e 100755 --- a/src/wrf/interp.py +++ b/src/wrf/interp.py @@ -4,9 +4,12 @@ from __future__ import (absolute_import, division, print_function, import numpy as np import numpy.ma as ma -from .extension import (interpz3d, interp2dxy, interp1d, - smooth2d, monotonic, vintrp, computevertcross, - computeinterpline) +# from .extension import (interpz3d, interp2dxy, interp1d, +# smooth2d, monotonic, vintrp, computevertcross, +# computeinterpline) + +from .extension import (_interpz3d, _interp2dxy, _interp1d, _vertcross, + _interpline, _smooth2d, monotonic, vintrp) from .metadecorators import set_interp_metadata from .util import extract_vars, is_staggered @@ -32,7 +35,7 @@ def interplevel(field3d, z, desiredloc, missingval=Constants.DEFAULT_FILL, missingval - the missing data value (which will be masked on return) """ - r1 = interpz3d(field3d, z, desiredloc, missingval) + r1 = _interpz3d(field3d, z, desiredloc, missingval) masked_r1 = ma.masked_values (r1, missingval) return masked_r1 @@ -64,7 +67,7 @@ def vertcross(field3d, z, missingval=Constants.DEFAULT_FILL, xy, var2dz, z_var2d = get_xy_z_params(z, pivot_point, angle, start_point, end_point) - res = computevertcross(field3d, xy, var2dz, z_var2d, missingval) + res = _vertcross(field3d, xy, var2dz, z_var2d, missingval) return ma.masked_values(res, missingval) @@ -89,7 +92,7 @@ def interpline(field2d, pivot_point=None, except (KeyError, TypeError): xy = get_xy(field2d, pivot_point, angle, start_point, end_point) - return computeinterpline(field2d, xy) + return _interpline(field2d, xy) @set_interp_metadata("vinterp") @@ -174,9 +177,8 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, ht_agl = get_height(wrfnc, timeidx, msl=False, units="m", method=method, squeeze=squeeze, cache=cache) - smsfp = smooth2d(sfp, 3) - - # Vertical coordinate type + smsfp = _smooth2d(sfp, 3) + vcor = 0 if vert_coord in ("pressure", "pres", "p"): @@ -244,17 +246,17 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, # TODO: Rename after the extensions are renamed @set_interp_metadata("horiz") def wrap_interpz3d(field3d, z, desiredloc, missingval, meta=True): - return interpz3d(field3d, z, desiredloc, missingval) + return _interpz3d(field3d, z, desiredloc, missingval) @set_interp_metadata("2dxy") def wrap_interp2dxy(field3d, xy, meta=True): - return interp2dxy(field3d, xy) + return _interp2dxy(field3d, xy) @set_interp_metadata("1d") def wrap_interp1d(v_in, z_in, z_out, missingval, meta=True): - return interp1d(v_in, z_in, z_out, missingval) + return _interp1d(v_in, z_in, z_out, missingval) diff --git a/src/wrf/interputils.py b/src/wrf/interputils.py index 50ab620..5bc91e2 100644 --- a/src/wrf/interputils.py +++ b/src/wrf/interputils.py @@ -5,7 +5,7 @@ from math import floor, ceil import numpy as np -from .extension import interp2dxy +from .extension import _interp2dxy from .util import py3range __all__ = ["to_positive_idxs", "calc_xy", "get_xy_z_params", "get_xy"] @@ -131,7 +131,7 @@ def get_xy_z_params(z, pivot_point=None, angle=None, xy = get_xy(z, pivot_point, angle, start_point, end_point) # Interp z - var2dz = interp2dxy(z, xy) + var2dz = _interp2dxy(z, xy) extra_dim_num = z.ndim - 3 idx1 = tuple([0]*extra_dim_num + [0,0]) diff --git a/src/wrf/omega.py b/src/wrf/omega.py index ae463e9..ec4de71 100755 --- a/src/wrf/omega.py +++ b/src/wrf/omega.py @@ -3,7 +3,8 @@ from __future__ import (absolute_import, division, print_function, from .constants import Constants from .destag import destagger -from .extension import computeomega,computetk +#from .extension import computeomega,computetk +from .extension import computeomega, _tk from .util import extract_vars from .metadecorators import copy_and_set_metadata @@ -26,7 +27,7 @@ def get_omega(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, wa = destagger(w, -3) full_t = t + Constants.T_BASE full_p = p + pb - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) omega = computeomega(qv,tk,wa,full_p) diff --git a/src/wrf/pw.py b/src/wrf/pw.py index cd7388e..e7f2fe1 100755 --- a/src/wrf/pw.py +++ b/src/wrf/pw.py @@ -1,7 +1,8 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -from .extension import computepw,computetv,computetk +#from .extension import computepw,computetv,computetk +from .extension import computepw,computetv, _tk from .constants import Constants from .util import extract_vars from .metadecorators import copy_and_set_metadata @@ -31,7 +32,7 @@ def get_pw(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, ht = (ph + phb)/Constants.G full_t = t + Constants.T_BASE - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) tv = computetv(tk, qv) return computepw(full_p, tv, qv, ht) diff --git a/src/wrf/rh.py b/src/wrf/rh.py index 1629a19..9b39a14 100755 --- a/src/wrf/rh.py +++ b/src/wrf/rh.py @@ -2,7 +2,8 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) from .constants import Constants -from .extension import computerh, computetk +#from .extension import computerh, computetk +from .extension import _rh, _tk from .util import extract_vars from .metadecorators import copy_and_set_metadata @@ -24,8 +25,8 @@ def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, full_t = t + Constants.T_BASE full_p = p + pb qvapor[qvapor < 0] = 0 - tk = computetk(full_p, full_t) - rh = computerh(qvapor, full_p, tk) + tk = _tk(full_p, full_t) + rh = _rh(qvapor, full_p, tk) return rh @@ -42,7 +43,7 @@ def get_rh_2m(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, q2 = ncvars["Q2"] q2[q2 < 0] = 0 - rh = computerh(q2, psfc, t2) + rh = _rh(q2, psfc, t2) return rh diff --git a/src/wrf/slp.py b/src/wrf/slp.py index cea61a9..3b9865f 100755 --- a/src/wrf/slp.py +++ b/src/wrf/slp.py @@ -1,7 +1,8 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -from .extension import computeslp, computetk +#from .extension import computeslp, computetk +from .extension import _slp, _tk from .constants import Constants from .destag import destagger from .decorators import convert_units @@ -37,8 +38,8 @@ def get_slp(wrfnc, timeidx=0, method="cat", squeeze=True, destag_ph = destagger(full_ph, -3) - tk = computetk(full_p, full_t) - slp = computeslp(destag_ph, tk, full_p, qvapor) + tk = _tk(full_p, full_t) + slp = _slp(destag_ph, tk, full_p, qvapor) return slp diff --git a/src/wrf/temp.py b/src/wrf/temp.py index f513c7f..f02eda3 100755 --- a/src/wrf/temp.py +++ b/src/wrf/temp.py @@ -2,7 +2,8 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) from .constants import Constants -from .extension import computetk, computeeth, computetv, computewetbulb +#from .extension import computetk, computeeth, computetv, computewetbulb +from .extension import _tk, computeeth, computetv, computewetbulb from .decorators import convert_units from .metadecorators import copy_and_set_metadata from .util import extract_vars @@ -42,7 +43,7 @@ def get_temp(wrfnc, timeidx=0, method="cat", squeeze=True, full_t = t + Constants.T_BASE full_p = p + pb - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) return tk @@ -64,7 +65,7 @@ def get_eth(wrfnc, timeidx=0, method="cat", squeeze=True, full_t = t + Constants.T_BASE full_p = p + pb - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) eth = computeeth(qv, tk, full_p) @@ -89,7 +90,7 @@ def get_tv(wrfnc, timeidx=0, method="cat", squeeze=True, full_t = t + Constants.T_BASE full_p = p + pb - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) tv = computetv(tk,qv) @@ -114,7 +115,7 @@ def get_tw(wrfnc, timeidx=0, method="cat", squeeze=True, full_t = t + Constants.T_BASE full_p = p + pb - tk = computetk(full_p, full_t) + tk = _tk(full_p, full_t) tw = computewetbulb(full_p,tk,qv) return tw diff --git a/src/wrf/util.py b/src/wrf/util.py index 6a644ae..2d17fa3 100644 --- a/src/wrf/util.py +++ b/src/wrf/util.py @@ -278,6 +278,27 @@ class combine_with(object): new_coords.update(self.new_coords) return new_dims, new_coords + +# This should look like: +# [(0, (-3,-2)), # variable 1 +# (1, -1)] # variable 2 +class combine_dims(object): + + def __init__(self, pairs): + self.pairs = pairs + + def __call__(self, *args): + result = [] + for pair in self.pairs: + var = args[pair[0]] + dimidxs = pair[1] + if isinstance(dimidxs, Iterable): + for dimidx in dimidxs: + result.append(var.shape[dimidx]) + else: + result.append(var.shape[dimidxs]) + + return tuple(result) def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): @@ -1291,8 +1312,9 @@ def _args_to_list2(func, args, kwargs): # Build the full tuple with defaults filled in outargs = [None]*len(argspec.args) - for i,default in enumerate(argspec.defaults[::-1], 1): - outargs[-i] = default + if argspec.defaults is not None: + for i,default in enumerate(argspec.defaults[::-1], 1): + outargs[-i] = default # Add the supplied args for i,arg in enumerate(args): diff --git a/src/wrf/uvdecorator.py b/src/wrf/uvdecorator.py index 9b1b664..eff190a 100644 --- a/src/wrf/uvdecorator.py +++ b/src/wrf/uvdecorator.py @@ -6,12 +6,15 @@ import numpy as np import wrapt from .destag import destagger -from .util import iter_left_indexes, py3range +from .util import iter_left_indexes, py3range, npvalues +from .config import xarray_enabled +from .constants import Constants + +if xarray_enabled(): + from xarray import DataArray __all__ = ["uvmet_left_iter"] -# Placed in separate module to resolve a circular dependency with destagger -# module def uvmet_left_iter(): """Decorator to handle iterating over leftmost dimensions when using @@ -53,13 +56,14 @@ def uvmet_left_iter(): return wrapped(u, v, lat, lon, cen_long, cone) # Start by getting the left-most 'extra' dims - outdims = [u.shape[x] for x in py3range(extra_dim_num)] + outdims = u.shape[0:extra_dim_num] extra_dims = list(outdims) # Copy the left-most dims for iteration # Append the right-most dimensions outdims += [2] # For u/v components - outdims += [u.shape[x] for x in py3range(-num_right_dims,0,1)] + #outdims += [u.shape[x] for x in py3range(-num_right_dims,0,1)] + outdims += list(u.shape[-num_right_dims:]) output = np.empty(outdims, u.dtype) @@ -75,16 +79,132 @@ def uvmet_left_iter(): new_lon = lon[left_and_slice_idxs] # Call the numerical routine - res = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone) + result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone) # Note: The 2D version will return a 3D array with a 1 length # dimension. Numpy is unable to broadcast this without # sqeezing first. - res = np.squeeze(res) + result = np.squeeze(result) - output[left_and_slice_idxs] = res[:] + output[left_and_slice_idxs] = result[:] return output return func_wrapper +def uvmet_left_iter_nocopy(alg_dtype=np.float64): + """Decorator to handle iterating over leftmost dimensions when using + multiple files and/or multiple times with the uvmet product. + + """ + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + u = args[0] + v = args[1] + lat = args[2] + lon = args[3] + cen_long = args[4] + cone = args[5] + + orig_dtype = u.dtype + + if u.ndim == lat.ndim: + num_right_dims = 2 + is_3d = False + else: + num_right_dims = 3 + is_3d = True + + has_missing = False + u_arr = u + if isinstance(u, DataArray): + u_arr = npvalues(u) + + v_arr = v + if isinstance(v, DataArray): + v_arr = npvalues(v) + + umissing = Constants.DEFAULT_FILL + if isinstance(u_arr, np.ma.MaskedArray): + has_missing = True + umissing = u_arr.fill_value + + vmissing = Constants.DEFAULT_FILL + if isinstance(v_arr, np.ma.MaskedArray): + has_missing = True + vmissing = v_arr.fill_value + + uvmetmissing = umissing + + is_stag = False + if (u.shape[-1] != lat.shape[-1] or u.shape[-2] != lat.shape[-2]): + is_stag = True + # Sanity check + if (v.shape[-1] == lat.shape[-1] or v.shape[-2] == lat.shape[-2]): + raise ValueError("u is staggered but v is not") + + if (v.shape[-1] != lat.shape[-1] or v.shape[-2] != lat.shape[-2]): + is_stag = True + # Sanity check + if (u.shape[-1] == lat.shape[-1] or u.shape[-2] == lat.shape[-2]): + raise ValueError("v is staggered but u is not") + + if is_3d: + extra_dim_num = u.ndim - 3 + else: + extra_dim_num = u.ndim - 2 + + + # No special left side iteration, return the function result + if (extra_dim_num == 0): + return wrapped(u, v, lat, lon, cen_long, cone, isstag=is_stag, + has_missing=has_missing, umissing=umissing, + vmissing=vmissing, uvmetmissing=uvmetmissing) + + # Start by getting the left-most 'extra' dims + outdims = u.shape[0:extra_dim_num] + extra_dims = list(outdims) # Copy the left-most dims for iteration + + # Append the right-most dimensions + if not is_3d: + outdims += (2,1) # Fortran routine needs 3 dimensions + else: + outdims += (2,) + + #outdims += [u.shape[x] for x in py3range(-num_right_dims,0,1)] + outdims += u.shape[-num_right_dims:] + + output = np.empty(outdims, alg_dtype) + + for left_idxs in iter_left_indexes(extra_dims): + + left_and_slice_idxs = left_idxs + (slice(None),) + + new_u = u[left_and_slice_idxs] + new_v = v[left_and_slice_idxs] + new_lat = lat[left_and_slice_idxs] + new_lon = lon[left_and_slice_idxs] + outview = output[left_and_slice_idxs] + + # Call the numerical routine + result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone, + isstag=is_stag, has_missing=has_missing, + umissing=umissing, vmissing=vmissing, + uvmetmissing=uvmetmissing, outview=outview) + + # Make sure the result is the same data as what got passed in + # Can delete this once everything works + if (result.__array_interface__["data"][0] != + outview.__array_interface__["data"][0]): + raise RuntimeError("output array was copied") + + + output = output.astype(orig_dtype) + + if has_missing: + output = np.ma.masked_values(output, uvmetmissing) + + return output.squeeze() + + return func_wrapper + diff --git a/src/wrf/uvmet.py b/src/wrf/uvmet.py index d9d7e31..3502d96 100755 --- a/src/wrf/uvmet.py +++ b/src/wrf/uvmet.py @@ -5,7 +5,8 @@ from math import fabs, log, tan, sin, cos import numpy as np -from .extension import computeuvmet +#from .extension import computeuvmet +from .extension import _uvmet from .destag import destagger from .constants import Constants from .wind import _calc_wspd_wdir @@ -120,9 +121,9 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, else: cone = 1 - res = computeuvmet(u, v ,lat, lon, cen_lon, cone) + result = _uvmet(u, v, lat, lon, cen_lon, cone) - return res + return result @set_wind_metadata(copy_varname=either("P", "PRES"), diff --git a/test/ipynb/WRF_python_demo.ipynb b/test/ipynb/WRF_python_demo.ipynb index 7e490e5..4c9370c 100644 --- a/test/ipynb/WRF_python_demo.ipynb +++ b/test/ipynb/WRF_python_demo.ipynb @@ -885,21 +885,21 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "Python 2", "language": "python", - "name": "python3" + "name": "python2" }, "language_info": { "codemirror_mode": { "name": "ipython", - "version": 3 + "version": 2 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.5.0rc4" + "pygments_lexer": "ipython2", + "version": "2.7.11" } }, "nbformat": 4, diff --git a/test/ipynb/nocopy_test.ipynb b/test/ipynb/nocopy_test.ipynb new file mode 100644 index 0000000..4243163 --- /dev/null +++ b/test/ipynb/nocopy_test.ipynb @@ -0,0 +1,234 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from wrf.extension import _slp" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from wrf import getvar\n", + "from netCDF4 import Dataset as nc\n", + "ncfile = nc(\"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d01_2005-08-28_00:00:00\")\n", + "b = getvar([ncfile,ncfile,ncfile], \"slp\", timeidx=None)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "print(b)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "b = getvar([ncfile,ncfile,ncfile], \"td\", None)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "print(b)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [ + "b = getvar([ncfile,ncfile,ncfile], \"tk\", None)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "print(b)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [ + "b = getvar([ncfile,ncfile,ncfile], \"rh\", None)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "print (b)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "# 500 MB Heights\n", + "from wrf import getvar, interplevel\n", + "\n", + "z = getvar([ncfile,ncfile,ncfile], \"z\", timeidx=None)\n", + "p = getvar([ncfile,ncfile,ncfile], \"pressure\", timeidx=None)\n", + "ht_500mb = interplevel(z, p, 500)\n", + "\n", + "print(ht_500mb)\n", + "del ht_500mb, z, p" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "# Pressure using pivot and angle\n", + "from wrf import getvar, vertcross\n", + "\n", + "z = getvar(ncfile, \"z\", timeidx=None)\n", + "p = getvar(ncfile, \"pressure\", timeidx=None)\n", + "pivot_point = (z.shape[-2] / 2, z.shape[-1] / 2) \n", + "angle = 90.0\n", + "\n", + "p_vert = vertcross(p, z, pivot_point=pivot_point, angle=angle)\n", + "print(p_vert)\n", + "del p_vert\n", + "\n", + "# Pressure using start_point and end_point\n", + "start_point = (z.shape[-2]/2, 0)\n", + "end_point = (z.shape[-2]/2, -1)\n", + "\n", + "p_vert = vertcross(p, z, start_point=start_point, end_point=end_point)\n", + "print(p_vert)\n", + "del p_vert, p, z\n" + ] + }, + { + "cell_type": "raw", + "metadata": {}, + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "# T2 using pivot and angle\n", + "from wrf import interpline, getvar\n", + "\n", + "t2 = getvar([ncfile,ncfile,ncfile], \"T2\", timeidx=None)\n", + "pivot_point = (t2.shape[-2] / 2, t2.shape[-1] / 2) \n", + "angle = 90.0\n", + "\n", + "t2_line = interpline(t2, pivot_point=pivot_point, angle=angle)\n", + "print(t2_line)\n", + "\n", + "del t2_line\n", + "\n", + "# T2 using start_point and end_point\n", + "start_point = (t2.shape[-2]/2, 0)\n", + "end_point = (t2.shape[-2]/2, -1)\n", + "\n", + "t2_line = interpline(t2, start_point=start_point, end_point=end_point)\n", + "print(t2_line)\n", + "\n", + "del t2_line, t2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from wrf import getvar\n", + "from netCDF4 import Dataset as nc\n", + "lambertnc = nc(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\")\n", + "uvmet = getvar([lambertnc,lambertnc], \"uvmet\", timeidx=None)\n", + "print (uvmet)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 2", + "language": "python", + "name": "python2" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2", + "version": "2.7.11" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/test/utests.py b/test/utests.py index d0dbb2e..25f1986 100644 --- a/test/utests.py +++ b/test/utests.py @@ -275,7 +275,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, timeidx=timeidx, log_p=True) - tol = 0/100. + tol = .5/100. atol = 0.0001 field = n.squeeze(field) @@ -296,7 +296,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, timeidx=timeidx, log_p=True) - tol = 0/100. + tol = .5/100. atol = 0.0001 field = n.squeeze(field) @@ -318,6 +318,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, log_p=True) field = n.squeeze(field) + nt.assert_allclose(npvalues(field), fld_tk_pres, tol, atol) # Tk to geoht_msl diff --git a/test/viewtest.py b/test/viewtest.py new file mode 100644 index 0000000..3f672b9 --- /dev/null +++ b/test/viewtest.py @@ -0,0 +1,21 @@ +import numpy as np + +import wrf._wrffortran + +a = np.ones((3,3,3)) +b = np.zeros((3,3,3,3)) +errstat = np.array(0) +errmsg = np.zeros(512, "c") + + +for i in xrange(2): + outview = b[i,:] + outview = outview.T + q = wrf._wrffortran.testfunc(a,outview,errstat=errstat,errstr=errmsg) + q[1,1,1] = 100 + + +print errstat +print b +str_bytes = (bytes(c).decode("utf-8") for c in errmsg[:]) +print "".join(str_bytes).strip() \ No newline at end of file