TST: Adding unit tests and PEP8 fixes to storm_relative_velocity fun.

pyart/retrieve/srv.py

@@ -1,18 +1,20 @@
 """
-Calculation of storm-relative velocity from a radar object. Code written by 
+Calculation of storm-relative velocity from a radar object. Code written by
 Edward C. Wolff. Modifications for single-sweep files suggested by Leanne
 Blind.
 
 """
 
-import numpy as np
 import math
 
-from ..config import get_field_name 
+import numpy as np
+
+from ..config import get_field_name
 
 
 def storm_relative_velocity(
-        radar, direction=None, speed=None, field=None, u=None, v=None):
+    radar, direction=None, speed=None, field=None, u=None, v=None
+):
     """
     This function calculates storm-relative Doppler velocities.
 
@@ -21,17 +23,17 @@ def storm_relative_velocity(
     radar: Radar
         Radar object used.
     direction: float or string
-        Direction of the storm motion vector (where north equals 0 degrees). 
-        Accepts a float or a string with the abbreviation of a cardinal or 
+        Direction of the storm motion vector (where north equals 0 degrees).
+        Accepts a float or a string with the abbreviation of a cardinal or
         ordinal/intercardinal direction (for example: N, SE, etc.). If both
         speed/direction and u/v are specified, speed/direction will be used.
     speed: string
         Speed of the storm motion vector.
-        Units should be identical to those in the provided radar 
+        Units should be identical to those in the provided radar
         object. If both speed/direction and u/v are specified, speed/direction
         will be used.
     field: string, optional
-        Velocity field to use for storm-relative calculation. A value of None 
+        Velocity field to use for storm-relative calculation. A value of None
         will use the default field name as defined in the Py-ART configuration
         file.
     u: float, optional
@@ -50,65 +52,64 @@ def storm_relative_velocity(
     # Parse the field parameter
     if field is None:
         field = get_field_name("velocity")
-    
+
     # Obtain velocity data and copy the array
-    sr_data = radar.fields[field]['data'].copy()
+    sr_data = radar.fields[field]["data"].copy()
 
     # Specify cardinal directions that can be interpreted
     direction_dict = {
-    "N": 0,
-    "NE": 45,
-    "E": 90,
-    "SE": 135,
-    "S": 180,
-    "SW": 225,
-    "W": 270,
-    "NW": 315,
+        "N": 0,
+        "NE": 45,
+        "E": 90,
+        "SE": 135,
+        "S": 180,
+        "SW": 225,
+        "W": 270,
+        "NW": 315,
     }
 
     # Set the direction of the storm motion vector
     # When speed and direction are specified
-    if direction is not None:
+    if direction is not None and speed is not None:
         if isinstance(direction, int) or isinstance(direction, float):
             alpha = direction
         elif isinstance(direction, str):
             if direction in direction_dict.keys():
                 alpha = direction_dict[direction]
             else:
-                raise ValueError('Direction string must be cardinal/ordinal direction')
+                raise ValueError("Direction string must be cardinal/ordinal direction")
         else:
-            raise ValueError('Direction must be an integer, float, or string')
+            raise ValueError("Direction must be an integer, float, or string")
     # When u and v are specified
     elif u is not None:
         if v is not None:
-            speed = np.sqrt((u**2)+(v**2))
-            direction = 90-np.rad2deg(math.atan2(v/speed, u/speed))
+            speed = np.sqrt((u**2) + (v**2))
+            direction = 90 - np.rad2deg(math.atan2(v / speed, u / speed))
             if direction < 0:
-                direction = direction+360
+                direction = direction + 360
         else:
-            raise ValueError('Must specify both u and v components')
+            raise ValueError("Must specify both u and v components")
     else:
-        raise ValueError('Must specify either speed and direction or u and v')
-
+        raise ValueError("Must specify either speed and direction or u and v")
 
     # Calculates the storm relative velocities
     # If the radar file contains only one sweep (e.g. some research radars)
-    if len(radar.sweep_number['data']) == 1:
+    if len(radar.sweep_number["data"]) == 1:
         sweep = 0
         start, end = radar.get_start_end(sweep)
         angle_array = radar.get_azimuth(sweep=sweep)
         ray_array = np.arange(start, end, 1)
         for count, ray in enumerate(ray_array):
-            correction = speed*np.cos(np.deg2rad(alpha-angle_array[count]))
-            sr_data[ray] = radar.fields[field]['data'][ray]-correction  
+            correction = speed * np.cos(np.deg2rad(alpha - angle_array[count]))
+            sr_data[ray] = radar.fields[field]["data"][ray] - correction
     # If the radar file contains several sweeps, one volume scan (e.g. NEXRAD)
     else:
-        for sweep in radar.sweep_number['data']:
+        for sweep in radar.sweep_number["data"]:
             start, end = radar.get_start_end(sweep)
             angle_array = radar.get_azimuth(sweep=sweep)
-            ray_array = np.arange(start, end+1, 1)
+            ray_array = np.arange(start, end + 1, 1)
             for count, ray in enumerate(ray_array):
-                correction = speed*np.cos(np.deg2rad(alpha-angle_array[count]))
-                sr_data[ray] = radar.fields[field]['data'][ray]-correction
-                
-    return sr_data
+                correction = speed * np.cos(np.deg2rad(alpha - angle_array[count]))
+                sr_data[ray] = radar.fields[field]["data"][ray] - correction
+
+    return sr_data

tests/retrieve/test_srv.py

@@ -0,0 +1,30 @@
+""" Unit Tests for Py-ART's retrieve/srv.py module. """
+
+import numpy as np
+import pytest
+
+import pyart
+
+
+def test_storm_relative_velocity():
+    # Test all sweeps
+    radar = pyart.io.read_nexrad_archive(
+        "s3://noaa-nexrad-level2/2022/03/31/KDGX/KDGX20220331_012808_V06"
+    )
+    sr_data = pyart.retrieve.storm_relative_velocity(radar, direction="NE", speed=20.0)
+    test_data = [2.276456117630005, 4.776455879211426, 3.276456117630005]
+    np.testing.assert_allclose(sr_data[-1][0:3], test_data)
+
+    # Test one sweep
+    radar_sweep = radar.extract_sweeps([21])
+    sr_one_sweep = pyart.retrieve.storm_relative_velocity(
+        radar_sweep, direction="NE", speed=20.0
+    )
+    one_sweep_data = [0.1278250813484192, 4.6278252601623535, 4.6278252601623535]
+    np.testing.assert_allclose(sr_one_sweep[0][0:3], one_sweep_data)
+
+    # Test missing parameters
+    pytest.raises(ValueError, pyart.retrieve.storm_relative_velocity, radar)
+    pytest.raises(
+        ValueError, pyart.retrieve.storm_relative_velocity, radar, u=14.14, v=None
+    )