Changed StatsMixin to be a child of MathMixin, added argmin and argmax functions

Author: d-giles

src/lkdata/datacube.py

@@ -16,7 +16,6 @@
 from .dataseries import DataSeries, BoolSeries, BitwiseSeries
 from .mixins import (
     StatsMixin,
-    MathMixin,
     BoolMixin,
     BitwiseMixin,
     AggMixin,
@@ -526,7 +525,6 @@ def values(self):
 
 
 class DataCube(
-    MathMixin,
     StatsMixin,
     Cube,
 ):

src/lkdata/dataseries.py

@@ -10,7 +10,6 @@
 from .mixins import (
     AggMixin,
     ConvenienceMixins,
-    MathMixin,
     StatsMixin,
     BoolMixin,
     BitwiseMixin,
@@ -134,7 +133,7 @@ def stats_post_process(self, result, **kwargs):
             return result
 
 
-class DataSeries(MathMixin, StatsMixin, Series):
+class DataSeries(StatsMixin, Series):
     """
     pandas.Series-like object with uncertainty and lightkurve functionality.
 

src/lkdata/mixins.py

@@ -37,6 +37,8 @@
     "sum",
     "std",
     "var",
+    "argmin",
+    "argmax",
     "min",
     "max",
     "prod",
@@ -777,7 +779,7 @@ def uncertainty(self, value):
         self._uncertainty = uncertainty
 
 
-class StatsMixin:
+class StatsMixin(MathMixin):
     """Defines a mixin class which will let us postprocess all our pandas stats"""
 
     _stats_type = "data"
@@ -795,6 +797,11 @@ def _create_stats_method(self, method_name):
         def _method(*args, **kwargs):
             axis = kwargs.pop("axis", None)
             np_method = getattr(np, method_name)
+            if np_method in (np.argmin, np.argmax):
+                if axis is None:
+                    # returning unravelled index instead of flattened index
+                    return np.unravel_index(np_method(self.array), self.array.shape)
+                return np_method(self.to_numpy(), axis=axis)
             result, init_kwds = self._arithmetic(
                 np_method, operand=None, data_axis=axis, uncertainty_axis=axis, **kwargs
             )

src/lkdata/seriescollection.py

@@ -13,7 +13,6 @@
 from .mixins import (
     AggMixin,
     ConvenienceMixins,
-    MathMixin,
     StatsMixin,
     BoolMixin,
     BitwiseMixin,
@@ -217,7 +216,7 @@ def stats_post_process(self, result, **kwargs):
             return result
 
 
-class DataSeriesCollection(MathMixin, StatsMixin, SeriesCollection):
+class DataSeriesCollection(StatsMixin, SeriesCollection):
     _series_class = DataSeries
 
     def __init__(

tests/test_datacube.py

@@ -38,12 +38,28 @@ def test_setup():
 
     # Test overridden pandas methods return correct shapes
     # Data products return tuples for and uncertainty
-    for method_name in STATS_METHOD_NAMES:
+    methods = [
+        method for method in STATS_METHOD_NAMES if method not in ["argmin", "argmax"]
+    ]
+    for method_name in methods:
         assert getattr(df, method_name)(axis=0)[0].shape == (nrow, ncol)
         assert (
             getattr(df[:, aperture], method_name)(axis=0)[0].shape[0] == aperture.sum()
         )
 
+    # argmin and argmax should return a single 3D index for axis=None
+    assert len(df.argmin()) == 3
+    assert len(df.argmax()) == 3
+
+    # argmin and argmax operate on axis=0 or 1 only
+    assert df.argmin(axis=0).shape[0] == nrow * ncol
+    assert df.argmax(axis=0).shape[0] == nrow * ncol
+    assert df.argmin(axis=1).shape[0] == ntime
+    assert df.argmax(axis=1).shape[0] == ntime
+
+    with pytest.raises(np.AxisError, match="axis 2 is out of bounds"):
+        _ = df.argmin(axis=2)
+
 
 def test_bad_setup():
     # Actual data values should be irrelevant for these tests