From 55d86cfad19c89cf9a2d5d7bf457c74ecdd94bc8 Mon Sep 17 00:00:00 2001 From: Andrew Date: Sun, 12 May 2024 22:53:17 +0100 Subject: [PATCH 1/2] Avoid looping on large numpy arrays --- pint/facets/numpy/numpy_func.py | 4 ++++ pint/testsuite/test_numpy.py | 5 +++++ 2 files changed, 9 insertions(+) diff --git a/pint/facets/numpy/numpy_func.py b/pint/facets/numpy/numpy_func.py index 138414553..6bccd409f 100644 --- a/pint/facets/numpy/numpy_func.py +++ b/pint/facets/numpy/numpy_func.py @@ -52,6 +52,10 @@ def _is_sequence_with_quantity_elements(obj): ------- True if obj is a sequence and at least one element is a Quantity; False otherwise """ + if np is not None and isinstance(obj, np.ndarray) and not obj.dtype.hasobject: + # If obj is a numpy array, avoid looping on all elements + # if dtype does not have objects + return False return ( iterable(obj) and sized(obj) diff --git a/pint/testsuite/test_numpy.py b/pint/testsuite/test_numpy.py index 486102124..d7221fa40 100644 --- a/pint/testsuite/test_numpy.py +++ b/pint/testsuite/test_numpy.py @@ -288,6 +288,11 @@ def test_broadcast_arrays(self): result = np.broadcast_arrays(x, y, subok=True) helpers.assert_quantity_equal(result, expected) + helpers.assert_quantity_equal( + np.broadcast_arrays(self.q[:, 1], np.array([1, 2])), + [np.array([[2, 4], [2, 4]]) * self.ureg.m, np.array([1, 2])], + ) + def test_roll(self): helpers.assert_quantity_equal( np.roll(self.q, 1), [[4, 1], [2, 3]] * self.ureg.m From d6a76977cf18dc6080ccfafdddf1e4b1c0f8a343 Mon Sep 17 00:00:00 2001 From: Andrew Date: Sun, 12 May 2024 23:06:50 +0100 Subject: [PATCH 2/2] Avoid looping on large numpy arrays --- pint/testsuite/test_numpy.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/pint/testsuite/test_numpy.py b/pint/testsuite/test_numpy.py index d7221fa40..486102124 100644 --- a/pint/testsuite/test_numpy.py +++ b/pint/testsuite/test_numpy.py @@ -288,11 +288,6 @@ def test_broadcast_arrays(self): result = np.broadcast_arrays(x, y, subok=True) helpers.assert_quantity_equal(result, expected) - helpers.assert_quantity_equal( - np.broadcast_arrays(self.q[:, 1], np.array([1, 2])), - [np.array([[2, 4], [2, 4]]) * self.ureg.m, np.array([1, 2])], - ) - def test_roll(self): helpers.assert_quantity_equal( np.roll(self.q, 1), [[4, 1], [2, 3]] * self.ureg.m