Fixed bug in np.concatenate() for instances of `named_arrays.Abstra…

…ctUncertainScalarArray` and `named_arrays.AbstractVectorArray`.
sun-data · Sep 27, 2023 · 514f12f · 514f12f
1 parent 799e5da
commit 514f12f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 13 deletions.
diff --git a/named_arrays/_scalars/uncertainties/uncertainties_array_functions.py b/named_arrays/_scalars/uncertainties/uncertainties_array_functions.py
@@ -539,15 +539,8 @@ def array_function_stack_like(
     arrays_nominal = []
     arrays_distribution = []
 
-    if func is np.concatenate:
-        shape = na.shape_broadcasted(*arrays)
-        if axis in shape:
-            shape = {axis: shape[axis]}
-        else:
-            shape = {}
-        arrays = tuple(na.broadcast_to(a, a.shape | shape) for a in arrays)
-
     for array in arrays:
+        array = array.broadcasted
         if isinstance(array, na.AbstractArray):
             if isinstance(array, na.AbstractScalar):
                 if isinstance(array, na.AbstractUncertainScalarArray):

diff --git a/named_arrays/_vectors/vector_array_functions.py b/named_arrays/_vectors/vector_array_functions.py
@@ -490,10 +490,11 @@ def array_function_stack_like(
             vector_prototype = array
             break
 
-    arrays = [vector_prototype.type_explicit.from_scalar(a) if not isinstance(a, na.AbstractVectorArray) else a
-              for a in arrays]
-    shape = na.shape_broadcasted(*arrays)
-    arrays = [a.broadcast_to(shape) for a in arrays]
+    arrays = [
+        vector_prototype.type_explicit.from_scalar(a) if not isinstance(a, na.AbstractVectorArray) else a
+        for a in arrays
+    ]
+    arrays = [a.broadcasted for a in arrays]
 
     components_arrays = [a.components for a in arrays]
 
@@ -505,7 +506,7 @@ def array_function_stack_like(
     components_result = dict()
     for c in components_arrays[0]:
         components_result[c] = func(
-            arrays=[components[c] for components in components_arrays],
+            [components[c] for components in components_arrays],
             axis=axis,
             out=components_out[c],
             dtype=dtype,