diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 15dd545a195..f833e6adf76 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -5,7 +5,6 @@ + ... ### New Features -+ `pm.math.cartesian` can now handle inputs that are themselves >1D (see [#4482](https://github.com/pymc-devs/pymc3/pull/4482)). + ... ### Maintenance diff --git a/pymc3/math.py b/pymc3/math.py index fc176b4185d..aff54d13b71 100644 --- a/pymc3/math.py +++ b/pymc3/math.py @@ -101,17 +101,11 @@ def cartesian(*arrays): Parameters ---------- - arrays: N-D array-like - N-D arrays where earlier arrays loop more slowly than later ones + arrays: 1D array-like + 1D arrays where earlier arrays loop more slowly than later ones """ N = len(arrays) - arrays_np = [np.asarray(x) for x in arrays] - arrays_2d = [x[:, None] if np.asarray(x).ndim == 1 else x for x in arrays_np] - arrays_integer = [np.arange(len(x)) for x in arrays_2d] - product_integers = np.stack(np.meshgrid(*arrays_integer, indexing="ij"), -1).reshape(-1, N) - return np.concatenate( - [array[product_integers[:, i]] for i, array in enumerate(arrays_2d)], axis=-1 - ) + return np.stack(np.meshgrid(*arrays, indexing="ij"), -1).reshape(-1, N) def kron_matrix_op(krons, m, op): diff --git a/pymc3/tests/test_math.py b/pymc3/tests/test_math.py index 5e73fb9c0f3..b31319021fd 100644 --- a/pymc3/tests/test_math.py +++ b/pymc3/tests/test_math.py @@ -71,24 +71,7 @@ def test_cartesian(): ] ) auto_cart = cartesian(a, b, c) - np.testing.assert_array_equal(manual_cartesian, auto_cart) - - -def test_cartesian_2d(): - np.random.seed(1) - a = [[1, 2], [3, 4]] - b = [5, 6] - c = [0] - manual_cartesian = np.array( - [ - [1, 2, 5, 0], - [1, 2, 6, 0], - [3, 4, 5, 0], - [3, 4, 6, 0], - ] - ) - auto_cart = cartesian(a, b, c) - np.testing.assert_array_equal(manual_cartesian, auto_cart) + np.testing.assert_array_almost_equal(manual_cartesian, auto_cart) def test_kron_dot():