From 4ebce59240a20b5a87a4b39709c743253528a163 Mon Sep 17 00:00:00 2001
From: Bill Ladwig <ladwig@ucar.edu>
Date: Fri, 12 Oct 2018 14:27:10 -0600
Subject: [PATCH] Fix storm relative helicity in southern hemisphere.

Latitude is now a required input when using data that has any values in the
southern hemisphere.
---
 fortran/wrf_relhl.f90  | 11 +++++++++--
 src/wrf/computation.py | 14 ++++++++++++--
 src/wrf/extension.py   |  9 +++++----
 src/wrf/g_helicity.py  |  6 +++++-
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/fortran/wrf_relhl.f90 b/fortran/wrf_relhl.f90
index b3e562f..aa9d3f4 100644
--- a/fortran/wrf_relhl.f90
+++ b/fortran/wrf_relhl.f90
@@ -31,7 +31,7 @@
 !   ***************************************************************
 
 ! NCLFORTSTART
-SUBROUTINE DCALRELHL(u, v, ght, ter, top, sreh, miy, mjx, mkzh)
+SUBROUTINE DCALRELHL(u, v, ght, ter, lat, top, sreh, miy, mjx, mkzh)
     USE wrf_constants, ONLY : PI, RAD_PER_DEG, DEG_PER_RAD
 
     IMPLICIT NONE
@@ -43,6 +43,7 @@ SUBROUTINE DCALRELHL(u, v, ght, ter, top, sreh, miy, mjx, mkzh)
     REAL(KIND=8), DIMENSION(miy,mjx,mkzh), INTENT(IN) :: u, v, ght
     REAL(KIND=8), INTENT(IN) :: top
     REAL(KIND=8), DIMENSION(miy,mjx), INTENT(IN) :: ter
+    REAL(KIND=8), DIMENSION(miy,mjx), INTENT(IN) :: lat
     REAL(KIND=8), DIMENSION(miy,mjx), INTENT(OUT) :: sreh
 
 ! NCLEND
@@ -100,7 +101,13 @@ SUBROUTINE DCALRELHL(u, v, ght, ter, top, sreh, miy, mjx, mkzh)
             ENDIF
 
             bsp = 0.75D0*asp
-            bdr = adr + 30.D0
+
+            IF (lat(i,j) .GE. 0) THEN ! Northern hemisphern
+                bdr = adr + 30.D0
+            ELSE ! Southern hemisphere
+                bdr = adr - 30.D0
+            END IF
+
             IF (bdr .GT. 360.D0) THEN
                 bdr = bdr - 360.D0
             ENDIF
diff --git a/src/wrf/computation.py b/src/wrf/computation.py
index 68ee1cf..53f3cf1 100644
--- a/src/wrf/computation.py
+++ b/src/wrf/computation.py
@@ -1298,7 +1298,7 @@ def dbz(pres, tkel, qv, qr, qs=None, qg=None, use_varint=False,
 
 @set_alg_metadata(2, "terrain", units="m2 s-2",
                   description="storm relative helicity")
-def srhel(u, v, height, terrain, top=3000.0, meta=True):
+def srhel(u, v, height, terrain, lats=None, top=3000.0, meta=True):
     """Return the storm relative helicity.
     
     This function calculates storm relative helicity from WRF ARW output. 
@@ -1342,6 +1342,11 @@ def srhel(u, v, height, terrain, top=3000.0, meta=True):
                 dimension names to the output.  Otherwise, default names will
                 be used. 
                 
+        lats (:class:`xarray.DataArray` or :class:`numpy.ndarray`, optional): 
+            Array of latitudes. This is required if any (or all) of your 
+            domain is in the southern hemisphere. If not provided, the 
+            northern hemisphere is assumed. Default is None.
+                
         top (:obj:`float`):  The height of the layer below which helicity is 
             calculated (meters above ground level).
 
@@ -1373,7 +1378,12 @@ def srhel(u, v, height, terrain, top=3000.0, meta=True):
     _v = np.ascontiguousarray(v[...,::-1,:,:])
     _height = np.ascontiguousarray(height[...,::-1,:,:])
     
-    return _srhel(_u, _v, _height, terrain, top)
+    if lats is None:
+        _lats = np.ones_like(terrain)
+    else:
+        _lats = lats
+    
+    return _srhel(_u, _v, _height, terrain, _lats, top)
 
 
 @set_alg_metadata(2, "u", refvarndims=3, units="m2 s-2",
diff --git a/src/wrf/extension.py b/src/wrf/extension.py
index 0c4355d..070a4b5 100755
--- a/src/wrf/extension.py
+++ b/src/wrf/extension.py
@@ -486,11 +486,11 @@ def _wetbulb(p, tk, qv,  psafile=psafilepath(), outview=None):
     return result
 
 
-@check_args(0, 3, (3,3,3,2))                         
-@left_iteration(3, 2, ref_var_idx=0, ignore_args=(4,))
-@cast_type(arg_idxs=(0,1,2,3))
+@check_args(0, 3, (3,3,3,2,2))                         
+@left_iteration(3, 2, ref_var_idx=0, ignore_args=(5,))
+@cast_type(arg_idxs=(0,1,2,3,4))
 @extract_and_transpose()
-def _srhel(u, v, z, ter, top, outview=None):
+def _srhel(u, v, z, ter, lats, top, outview=None):
     """Wrapper for dcalrelhl.
     
     Located in wrf_relhl.f90.
@@ -503,6 +503,7 @@ def _srhel(u, v, z, ter, top, outview=None):
                        v, 
                        z, 
                        ter, 
+                       lats,
                        top,
                        outview)
     
diff --git a/src/wrf/g_helicity.py b/src/wrf/g_helicity.py
index 36464f9..5ae17ab 100755
--- a/src/wrf/g_helicity.py
+++ b/src/wrf/g_helicity.py
@@ -7,6 +7,7 @@ from .extension import _srhel, _udhel
 from .destag import destagger
 from .util import extract_vars, extract_global_attrs, either
 from .metadecorators import copy_and_set_metadata
+from .g_latlon import get_lat
 
 @copy_and_set_metadata(copy_varname="HGT", name="srh", 
                        description="storm relative helicity",
@@ -72,6 +73,9 @@ def get_srh(wrfin, timeidx=0, method="cat", squeeze=True,
     """
     # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh)
     
+    lats = get_lat(wrfin, timeidx, method, squeeze, 
+                   cache, meta=False, _key=_key, stagger=None)
+    
     ncvars = extract_vars(wrfin, timeidx, ("HGT", "PH", "PHB"),
                           method, squeeze, cache, meta=False,
                           _key=_key)
@@ -101,7 +105,7 @@ def get_srh(wrfin, timeidx=0, method="cat", squeeze=True,
     v1 = np.ascontiguousarray(v[...,::-1,:,:])
     z1 = np.ascontiguousarray(z[...,::-1,:,:])
     
-    srh = _srhel(u1, v1, z1, ter, top)
+    srh = _srhel(u1, v1, z1, ter, lats, top)
     
     return srh