angle_to_direction handle ndarray

src/metpy/calc/tools.py

@@ -1845,8 +1845,11 @@
     else:
         scalar = False
 
+    np_input_angle = np.array(input_angle).astype(float)
+    origshape = np_input_angle.shape
+    ndarray = len(origshape) > 1
     # clean any numeric strings, negatives, and None does not handle strings with alphabet
-    input_angle = units.Quantity(np.array(input_angle).astype(float), origin_units)
+    input_angle = units.Quantity(np_input_angle, origin_units)
     input_angle[input_angle < 0] = np.nan
 
     # Normalize between 0 - 360
@@ -1862,8 +1865,10 @@
         err_msg = 'Level of complexity cannot be less than 1 or greater than 3!'
         raise ValueError(err_msg)
 
-    angle_dict = {i * BASE_DEGREE_MULTIPLIER.m * nskip: dir_str
-                  for i, dir_str in enumerate(DIR_STRS[::nskip])}
+    angle_dict = {
+        i * BASE_DEGREE_MULTIPLIER.m * nskip: dir_str
+        for i, dir_str in enumerate(DIR_STRS[::nskip])
+    }
     angle_dict[MAX_DEGREE_ANGLE.m] = 'N'  # handle edge case of 360.
     angle_dict[UND_ANGLE] = UND
 
@@ -1879,18 +1884,25 @@
     # ['N', 'N', 'NE', 'NE', 'E', 'E', 'SE', 'SE',
     #  'S', 'S', 'SW', 'SW', 'W', 'W', 'NW', 'NW']
 
-    multiplier = np.round(
-        (norm_angles / BASE_DEGREE_MULTIPLIER / nskip) - 0.001).m
-    round_angles = (multiplier * BASE_DEGREE_MULTIPLIER.m * nskip)
+    multiplier = np.round((norm_angles / BASE_DEGREE_MULTIPLIER / nskip) - 0.001).m
+    round_angles = multiplier * BASE_DEGREE_MULTIPLIER.m * nskip
     round_angles[np.where(np.isnan(round_angles))] = UND_ANGLE
-
-    dir_str_arr = itemgetter(*round_angles)(angle_dict)  # for array
-    if not full:
-        return dir_str_arr
-
-    dir_str_arr = ','.join(dir_str_arr)
-    dir_str_arr = _unabbreviate_direction(dir_str_arr)
-    return dir_str_arr.replace(',', ' ') if scalar else dir_str_arr.split(',')
+    if ndarray:
+        round_angles = round_angles.flatten()
+    dir_str_arr = itemgetter(*round_angles)(angle_dict)  # returns str or tuple
+    if full:
+        dir_str_arr = ','.join(dir_str_arr)
+        dir_str_arr = _unabbreviate_direction(dir_str_arr)
+        dir_str_arr = dir_str_arr.split(',')
+        if scalar:
+            return dir_str_arr[0]
+        else:
+            return np.array(dir_str_arr).reshape(origshape)
+    else:
+        if scalar:
+            return dir_str_arr
+        else:
+            return np.array(dir_str_arr).reshape(origshape)
 
 
 def _unabbreviate_direction(abb_dir_str):

tests/calc/test_calc_tools.py

@@ -933,6 +933,14 @@ def test_angle_to_direction_level_1():
     assert_array_equal(output_dirs, expected_dirs)
 
 
+def test_angle_to_direction_ndarray():
+    """Test array of angles in degree with a 2d numpy array."""
+    expected_dirs = np.array([['E', 'W'], ['E', 'W']])
+    input_angle = np.array([[90, 270], [90, 270]])
+    output_dirs = angle_to_direction(input_angle, level=1)
+    assert_array_equal(output_dirs, expected_dirs)
+
+
 def test_azimuth_range_to_lat_lon():
     """Test conversion of azimuth and range to lat/lon grid."""
     az = [332.2403, 334.6765, 337.2528, 339.73846, 342.26257]