From 09117adecf4e2fe11aaef96825520e7379244d4e Mon Sep 17 00:00:00 2001 From: Sven Baars Date: Wed, 20 Mar 2024 17:19:11 +0100 Subject: [PATCH] quantities: Make mean() more generic. (#1039) --- src/amuse/test/suite/core_tests/test_quantities.py | 6 ++++++ src/amuse/units/quantities.py | 4 ++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/amuse/test/suite/core_tests/test_quantities.py b/src/amuse/test/suite/core_tests/test_quantities.py index 6d82c07836..25bf8ef984 100644 --- a/src/amuse/test/suite/core_tests/test_quantities.py +++ b/src/amuse/test/suite/core_tests/test_quantities.py @@ -415,6 +415,12 @@ def test34(self): x[0] = -1 self.assertEqual(a, [1, 2, 3, 4, 5] | units.m) + def test_mean(self): + x = [1.0, 2.0, 3.0] | si.kg + mask = [True, True, False] + self.assertEqual(x.mean(), 2.0 | si.kg) + self.assertEqual(x.mean(where=mask), 1.5 | si.kg) + class TestAdaptingVectorQuantities(amusetest.TestCase): diff --git a/src/amuse/units/quantities.py b/src/amuse/units/quantities.py index cd75f2a9cc..8bb0bd43ea 100644 --- a/src/amuse/units/quantities.py +++ b/src/amuse/units/quantities.py @@ -973,8 +973,8 @@ def transpose(self, axes=None): def T(self): return VectorQuantity(self.number.T, self.unit) - def mean(self, axis=None, dtype=None, out=None): - return new_quantity(self.number.mean(axis, dtype, out), self.unit) + def mean(self, *args, **kwargs): + return new_quantity(self.number.mean(*args, **kwargs), self.unit) def median(self, **kwargs): return new_quantity(numpy.median(self.number, **kwargs), self.unit)