diff --git a/pysindy/utils/axes.py b/pysindy/utils/axes.py index bad10d55c..ad0f79040 100644 --- a/pysindy/utils/axes.py +++ b/pysindy/utils/axes.py @@ -117,13 +117,16 @@ def decorator(func): @implements(np.concatenate) -def concatenate(arrays, axis=0): +def concatenate(arrays, axis=0, out=None, dtype=None, casting="same_kind"): parents = [np.asarray(obj) for obj in arrays] ax_list = [obj.__dict__ for obj in arrays if isinstance(obj, AxesArray)] for ax1, ax2 in zip(ax_list[:-1], ax_list[1:]): if ax1 != ax2: raise TypeError("Concatenating >1 AxesArray with incompatible axes") - return AxesArray(np.concatenate(parents, axis), axes=ax_list[0]) + result = np.concatenate(parents, axis, out=out, dtype=dtype, casting=casting) + if isinstance(out, AxesArray): + out.__dict__ = ax_list[0] + return AxesArray(result, axes=ax_list[0]) def comprehend_axes(x): diff --git a/test/utils/test_axes.py b/test/utils/test_axes.py index b1a38e6f4..e5d9a8385 100644 --- a/test/utils/test_axes.py +++ b/test/utils/test_axes.py @@ -7,6 +7,13 @@ from pysindy import AxesArray +def test_concat_out(): + arr = AxesArray(np.arange(3).reshape(1, 3), {"ax_a": 0, "ax_b": 1}) + arr_out = np.empty((2, 3)).view(AxesArray) + result = np.concatenate((arr, arr), axis=0, out=arr_out) + assert_equal(result, arr_out) + + def test_reduce_mean_noinf_recursion(): arr = AxesArray(np.array([[1]]), {}) np.mean(arr, axis=0)