Merge pull request #2204 from cphyc/fix/missing_argument_ustack

[BUGFIX] Fix missing argument in stack ufunc
yt-project · Mar 13, 2019 · 1019292 · 1019292
2 parents 81e885c + 9d7f299
commit 1019292
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/yt/units/tests/test_ytarray.py b/yt/units/tests/test_ytarray.py
@@ -1281,6 +1281,7 @@ def test_numpy_wrappers():
     union_answer = [1, 2, 3, 4, 5, 6]
     vstack_answer = [[2, 3, 4, 5, 6],
                      [7, 8, 9,10, 11]]
+    vstack_answer_last_axis = [[ 2, 7], [ 3, 8], [ 4, 9], [ 5, 10], [ 6, 11]]
 
     assert_array_equal(YTArray(catenate_answer, 'cm'),
                        uconcatenate((a1, a2)))
@@ -1302,6 +1303,9 @@ def test_numpy_wrappers():
     assert_array_equal(YTArray(vstack_answer, 'cm'), ustack([a2, a3]))
     assert_array_equal(vstack_answer, np.stack([a2, a3]))
 
+    assert_array_equal(YTArray(vstack_answer_last_axis, 'cm'), ustack([a2, a3], axis=-1))
+    assert_array_equal(vstack_answer_last_axis, np.stack([a2, a3], axis=-1))
+
 def test_dimensionless_conversion():
     a = YTQuantity(1, 'Zsun')
     b = a.in_units('Zsun')

diff --git a/yt/units/yt_array.py b/yt/units/yt_array.py
@@ -1693,7 +1693,7 @@ def ustack(arrs, axis=0):
     This is a wrapper around np.stack that preserves units.
 
     """
-    v = np.stack(arrs)
+    v = np.stack(arrs, axis=axis)
     v = validate_numpy_wrapper_units(v, arrs)
     return v