Updated wspd wdir to match NCL

fortran/wrf_wind.f90

@@ -1,23 +1,21 @@
 ! NCLFORTSTART
-SUBROUTINE DCOMPUTEWSPD(wspd, u, v, nx, ny)
+SUBROUTINE DCOMPUTEWSPD(wspd, u, v, n)
 
     IMPLICIT NONE
 
     !f2py threadsafe
     !f2py intent(in,out) :: wspd
 
-    INTEGER, INTENT(IN) :: nx, ny
-    REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: wspd
-    REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: u, v
+    INTEGER, INTENT(IN) :: n
+    REAL(KIND=8), DIMENSION(n), INTENT(OUT) :: wspd
+    REAL(KIND=8), DIMENSION(n), INTENT(IN) :: u, v
 ! NCLEND
 
-    INTEGER i, j
+    INTEGER i
 
-    !$OMP PARALLEL DO COLLAPSE(2) SCHEDULE(runtime)
-    DO j = 1,ny
-        DO i = 1,nx
-            wspd(i,j) = SQRT(u(i,j)*u(i,j) + v(i,j)*v(i,j))
-        END DO
+    !$OMP PARALLEL DO SCHEDULE(runtime)
+    DO i = 1,n
+        wspd(i) = SQRT(u(i)*u(i) + v(i)*v(i))
     END DO
     !$OMP END PARALLEL DO
 
@@ -25,7 +23,7 @@ END SUBROUTINE DCOMPUTEWSPD
 
 
 ! NCLFORTSTART
-SUBROUTINE DCOMPUTEWDIR(wdir, u, v, nx, ny)
+SUBROUTINE DCOMPUTEWDIR(wdir, u, v, n)
     USE wrf_constants, ONLY : DEG_PER_RAD
 
     IMPLICIT NONE
@@ -33,18 +31,16 @@ SUBROUTINE DCOMPUTEWDIR(wdir, u, v, nx, ny)
     !f2py threadsafe
     !f2py intent(in,out) :: wdir
 
-    INTEGER, INTENT(IN) :: nx, ny
-    REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: wdir
-    REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: u, v
+    INTEGER, INTENT(IN) :: n
+    REAL(KIND=8), DIMENSION(n), INTENT(OUT) :: wdir
+    REAL(KIND=8), DIMENSION(n), INTENT(IN) :: u, v
 ! NCLEND
 
-    INTEGER i, j
+    INTEGER i
 
-    !$OMP PARALLEL DO COLLAPSE(2) SCHEDULE(runtime)
-    DO j = 1,ny
-        DO i = 1,nx
-            wdir(i,j) = MOD(270.0 - ATAN2(v(i,j), u(i,j)) * DEG_PER_RAD, 360.)
-        END DO
+    !$OMP PARALLEL DO SCHEDULE(runtime)
+    DO i = 1,n
+        wdir(i) = MOD(270.0 - ATAN2(v(i), u(i)) * DEG_PER_RAD, 360.)
     END DO
     !$OMP END PARALLEL DO
 

src/wrf/extension.py

@@ -949,12 +949,15 @@ def _wspd(u, v, outview=None):
     Located in wrf_wind.f90.
 
     """
+    shape = u.shape
     if outview is None:
         outview = np.empty_like(u)
 
-    result = dcomputewspd(outview,
-                          u,
-                          v)
+    result = dcomputewspd(outview.ravel(order="A"),
+                          u.ravel(order="A"),
+                          v.ravel(order="A"))
+
+    result = np.reshape(result, shape, order="F")
 
     return result
 
@@ -969,12 +972,15 @@ def _wdir(u, v, outview=None):
     Located in wrf_wind.f90.
 
     """
+    shape = u.shape
     if outview is None:
         outview = np.empty_like(u)
 
-    result = dcomputewdir(outview,
-                          u,
-                          v)
+    result = dcomputewdir(outview.ravel(order="A"),
+                          u.ravel(order="A"),
+                          v.ravel(order="A"))
+
+    result = np.reshape(result, shape, order="F")
 
     return result