Skip to content

Commit

Permalink
Avoid looping on large numpy arrays
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgsavage committed May 12, 2024
1 parent 24dd237 commit 55d86cf
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 0 deletions.
4 changes: 4 additions & 0 deletions pint/facets/numpy/numpy_func.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions pint/testsuite/test_numpy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 55d86cf

Please sign in to comment.