quantities: Make mean() more generic. (amusecode#1039)

rieder · Mar 20, 2024 · 09117ad · 09117ad
1 parent 51443e9
commit 09117ad
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git 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
@@ -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)