Repository: yt-project/unyt
Pull request #225 (2022-03-30 18:42:12)

BUG: fix an issue where multiplying a 1-element unyt_array would return a unyt_quantity
This adresses one of the two bugs reported in #224 

~Based on top of #219 for stability, only the last two commits are exclusive to this branch~
edit: rebased on master !


Patch:
diff --git a/unyt/array.py b/unyt/array.py
index ac7085b..9ef5715 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -1897,8 +1897,10 @@ class unyt_array(np.ndarray):
             out_arr = np.array(out_arr, copy=False)
         elif ufunc in (modf, divmod_):
             out_arr = tuple((ret_class(o, unit) for o in out_arr))
-        elif out_arr.size == 1:
+        elif out_arr.shape == ():
             out_arr = unyt_quantity(np.asarray(out_arr), unit)
+        elif out_arr.size == 1:
+            out_arr = unyt_array(np.asarray(out_arr), unit)
         else:
             if ret_class is unyt_quantity:
                 # This happens if you do ndarray * unyt_quantity.
@@ -2144,6 +2146,17 @@ class unyt_quantity(unyt_array):
     def __round__(self):
         return type(self)(round(float(self)), self.units)
 
+    def reshape(self, shape, order="C"):
+        # this is necessary to support some numpy operations
+        # natively, like numpy.meshgrid, which internally performs
+        # reshaping, e.g., arr.reshape(1, -1), which doesn't affect the size,
+        # but does change the object's internal representation to a >0D array
+        # see https://github.com/yt-project/unyt/issues/224
+        if shape == () or shape is None:
+            return super().reshape(shape, order=order)
+        else:
+            return unyt_array(self).reshape(shape, order=order)
+
 
 def _validate_numpy_wrapper_units(v, arrs):
     if not any(isinstance(a, unyt_array) for a in arrs):


Test patch:
diff --git a/unyt/tests/test_unyt_array.py b/unyt/tests/test_unyt_array.py
index b8f5530..04bcea1 100644
--- a/unyt/tests/test_unyt_array.py
+++ b/unyt/tests/test_unyt_array.py
@@ -2592,3 +2592,44 @@ def test_invalid_unit_quantity_from_string(s):
         match="Could not find unit symbol '{}' in the provided symbols.".format(un_str),
     ):
         unyt_quantity.from_string(s)
+
+
+def test_constant_type():
+    # see https://github.com/yt-project/unyt/issues/224
+    a = [1] * cm
+    assert type(a) is unyt_array
+    b = 2 * a
+    assert type(b) is unyt_array
+
+
+def test_composite_meshgrid():
+    # see https://github.com/yt-project/unyt/issues/224
+    a = np.array(1)
+    # pure numpy call to illustrate that the problem
+    # is only with units
+    np.meshgrid(np.array([1, 2]), a)
+    np.meshgrid(np.array([1, 2]), a * m)
+
+
+@pytest.mark.parametrize(
+    "shape, expected_output_shape",
+    [
+        (1, (1,)),
+        ((1,), (1,)),
+        ((1, 1), (1, 1)),
+        ((1, -1), (1, 1)),
+    ],
+)
+def test_reshape_quantity_to_array(shape, expected_output_shape):
+    a = unyt_quantity(1, "m")
+    b = a.reshape(shape)
+    assert b.shape == expected_output_shape
+    assert type(b) is unyt_array
+
+
+@pytest.mark.parametrize("shape", ((), None))
+def test_reshape_quantity_noop(shape):
+    a = unyt_quantity(1, "m")
+    b = a.reshape(shape)
+    assert b.shape == a.shape == ()
+    assert type(b) is unyt_quantity