Merge pull request NCAS-CMS#648 from davidhassell/argmin

`Data.argmin` and `Field.argmin` methods

Author: davidhassell

Changelog.rst

@@ -3,6 +3,8 @@ version 3.15.1
 
 **2023-05-??**
 
+* New methods: `cf.Data.argmin`, `cf.Field.argmin`
+  (https://github.com/NCAS-CMS/cf-python/issues/577)
 * Fix bug when using the ``-d`` option to the `cfa` script 
   (https://github.com/NCAS-CMS/cf-python/issues/649)
 

cf/data/data.py

@@ -5850,6 +5850,10 @@ def argmax(self, axis=None, unravel=False):
         If the data index is returned as a `tuple` (see the *unravel*
         parameter) then all delayed operations are computed.
 
+        .. versionadded:: 3.0.0
+
+        .. seealso:: `argmin`
+
         :Parameters:
 
             axis: `int`, optional
@@ -5878,7 +5882,7 @@ def argmax(self, axis=None, unravel=False):
         >>> a = d.argmax()
         >>> a
         <CF Data(): 5>
-        >>> a.array
+        >>> print(a.array)
         5
 
         >>> index = d.argmax(unravel=True)
@@ -5903,7 +5907,7 @@ def argmax(self, axis=None, unravel=False):
         >>> print(d.array)
         [[0 1 2]
          [3 5 5]]
-        >>> d.argmax(1)
+        >>> d.argmax(axis=1)
         <CF Data(2): [2, 1]>
 
         """
@@ -5916,6 +5920,90 @@ def argmax(self, axis=None, unravel=False):
 
         return type(self)(a)
 
+    def argmin(self, axis=None, unravel=False):
+        """Return the indices of the minimum values along an axis.
+
+        If no axis is specified then the returned index locates the
+        minimum of the whole data.
+
+        In case of multiple occurrences of the minimum values, the
+        indices corresponding to the first occurrence are returned.
+
+        **Performance**
+
+        If the data index is returned as a `tuple` (see the *unravel*
+        parameter) then all delayed operations are computed.
+
+        .. versionadded:: 3.15.1
+
+        .. seealso:: `argmax`
+
+        :Parameters:
+
+            axis: `int`, optional
+                The specified axis over which to locate the minimum
+                values. By default the minimum over the flattened data
+                is located.
+
+            unravel: `bool`, optional
+                If True then when locating the minimum over the whole
+                data, return the location as an integer index for each
+                axis as a `tuple`. By default an index to the
+                flattened array is returned in this case. Ignored if
+                locating the minima over a subset of the axes.
+
+        :Returns:
+
+            `Data` or `tuple` of `int`
+                The location of the minimum, or minima.
+
+        **Examples**
+
+        >>> d = cf.Data(np.arange(5, -1, -1).reshape(2, 3))
+        >>> print(d.array)
+        [[5 4 3]
+         [2 1 0]]
+        >>> a = d.argmin()
+        >>> a
+        <CF Data(): 5>
+        >>> print(a.array)
+        5
+
+        >>> index = d.argmin(unravel=True)
+        >>> index
+        (1, 2)
+        >>> d[index]
+        <CF Data(1, 1): [[0]]>
+
+        >>> d.argmin(axis=0)
+        <CF Data(3): [1, 1, 1]>
+        >>> d.argmin(axis=1)
+        <CF Data(2): [2, 2]>
+
+        Only the location of the first occurrence is returned:
+
+        >>> d = cf.Data([4, 0, 2, 3, 0])
+        >>> d.argmin()
+        <CF Data(): 1>
+
+        >>> d = cf.Data(np.arange(5, -1, -1).reshape(2, 3))
+        >>> d[1, 1] = 0
+        >>> print(d.array)
+        [[5 4 3]
+         [2 0 0]]
+        >>> d.argmin(axis=1)
+        <CF Data(2): [2, 1]>
+
+        """
+        dx = self.to_dask_array()
+        a = dx.argmin(axis=axis)
+
+        if unravel and (axis is None or self.ndim <= 1):
+            # Return a multidimensional index tuple
+            return tuple(np.array(da.unravel_index(a, self.shape)))
+
+        return type(self)(a)
+
     @_inplace_enabled(default=False)
     def convert_reference_time(
         self,

cf/field.py

@@ -11750,7 +11750,9 @@ def argmax(self, axis=None, unravel=False):
         If the data index is returned as a `tuple` (see the *unravel*
         parameter) then all delayed operations are computed.
 
-        .. seealso:: `where`, `cf.Data.argmax`
+        .. versionadded:: 3.0.0
+
+        .. seealso:: `argmin`, `where`, `cf.Data.argmax`
 
         :Parameters:
 
@@ -11825,6 +11827,97 @@ def argmax(self, axis=None, unravel=False):
 
         return self.data.argmax(axis=axis, unravel=unravel)
 
+    def argmin(self, axis=None, unravel=False):
+        """Return the indices of the minimum values along an axis.
+
+        If no axis is specified then the returned index locates the
+        minimum of the whole data.
+
+        In case of multiple occurrences of the minimum values, the
+        indices corresponding to the first occurrence are returned.
+
+        **Performance**
+
+        If the data index is returned as a `tuple` (see the *unravel*
+        parameter) then all delayed operations are computed.
+
+        .. versionadded:: 3.15.1
+
+        .. seealso:: `argmax`, `where`, `cf.Data.argmin`
+
+        :Parameters:
+
+            axis: optional
+                Select the domain axis over which to locate the
+                minimum values, defined by the domain axis that would
+                be selected by passing the given *axis* to a call of
+                the field construct's `domain_axis` method. For
+                example, for a value of ``'X'``, the domain axis
+                construct returned by ``f.domain_axis('X')`` is
+                selected.
+
+                By default the minimum over the flattened data is
+                located.
+
+            unravel: `bool`, optional
+                If True then when locating the minimum over the whole
+                data, return the location as an integer index for each
+                axis as a `tuple`. By default an index to the
+                flattened array is returned in this case. Ignored if
+                locating the minima over a subset of the axes.
+
+        :Returns:
+
+            `Data` or `tuple` of `int`
+                The location of the minimum, or minima.
+
+        **Examples**
+
+        >>> f = cf.example_field(2)
+        >>> print(f)
+        Field: air_potential_temperature (ncvar%air_potential_temperature)
+        ------------------------------------------------------------------
+        Data            : air_potential_temperature(time(36), latitude(5), longitude(8)) K
+        Cell methods    : area: mean
+        Dimension coords: time(36) = [1959-12-16 12:00:00, ..., 1962-11-16 00:00:00]
+                        : latitude(5) = [-75.0, ..., 75.0] degrees_north
+                        : longitude(8) = [22.5, ..., 337.5] degrees_east
+                        : air_pressure(1) = [850.0] hPa
+
+        Find the T axis indices of the minimum at each X-Y location:
+
+        >>> i = f.argmin('T')
+        >>> print(i.array)
+        [[ 5 14  4 32 11 34  9 27]
+         [10  7  4 21 11 10  3 33]
+         [21 33  1 30 26  8 33  4]
+         [15 10 12 19 23 20 30 25]
+         [28 33 31 11 15 12  9  7]]
+
+        Find the coordinates of the global minimum value, showing that
+        it occurs on 1960-03-16 12:00:00 at location 292.5 degrees
+        east, -45.0 degrees north:
+
+        >>> g = f[f.argmin(unravel=True)]
+        >>> print(g)
+        Field: air_potential_temperature (ncvar%air_potential_temperature)
+        ------------------------------------------------------------------
+        Data            : air_potential_temperature(time(1), latitude(1), longitude(1)) K
+        Cell methods    : area: mean
+        Dimension coords: time(1) = [1960-03-16 12:00:00]
+                        : latitude(1) = [-45.0] degrees_north
+                        : longitude(1) = [292.5] degrees_east
+                        : air_pressure(1) = [850.0] hPa
+
+        See `cf.Data.argmin` for further examples.
+
+        """
+        if axis is not None:
+            axis = self.domain_axis(axis, key=True)
+            axis = self.get_data_axes().index(axis)
+
+        return self.data.argmin(axis=axis, unravel=unravel)
+
     @_deprecated_kwarg_check("i", version="3.0.0", removed_at="4.0.0")
     def squeeze(self, axes=None, inplace=False, i=False, **kwargs):
         """Remove size 1 axes from the data.

cf/test/test_Data.py

@@ -2476,21 +2476,21 @@ def test_Data__int__(self):
 
     def test_Data_argmax(self):
         """Test the `argmax` Data method."""
-        d = cf.Data(np.arange(120).reshape(4, 5, 6))
+        d = cf.Data(np.arange(24).reshape(2, 3, 4))
 
-        self.assertEqual(d.argmax().array, 119)
+        self.assertEqual(d.argmax().array, 23)
 
         index = d.argmax(unravel=True)
-        self.assertEqual(index, (3, 4, 5))
-        self.assertEqual(d[index].array, 119)
+        self.assertEqual(index, (1, 2, 3))
+        self.assertEqual(d[index].array, 23)
 
         e = d.argmax(axis=1)
-        self.assertEqual(e.shape, (4, 6))
+        self.assertEqual(e.shape, (2, 4))
         self.assertTrue(
-            e.equals(cf.Data.full(shape=(4, 6), fill_value=4, dtype=int))
+            e.equals(cf.Data.full(shape=(2, 4), fill_value=2, dtype=int))
         )
 
-        self.assertEqual(d[d.argmax(unravel=True)].array, 119)
+        self.assertEqual(d[d.argmax(unravel=True)].array, 23)
 
         d = cf.Data([0, 4, 2, 3, 4])
         self.assertEqual(d.argmax().array, 1)
@@ -2499,6 +2499,31 @@ def test_Data_argmax(self):
         with self.assertRaises(ValueError):
             d.argmax(axis=d.ndim)
 
+    def test_Data_argmin(self):
+        """Test the `argmin` Data method."""
+        d = cf.Data(np.arange(23, -1, -1).reshape(2, 3, 4))
+
+        self.assertEqual(d.argmin().array, 23)
+
+        index = d.argmin(unravel=True)
+        self.assertEqual(index, (1, 2, 3))
+        self.assertEqual(d[index].array, 0)
+
+        e = d.argmin(axis=1)
+        self.assertEqual(e.shape, (2, 4))
+        self.assertTrue(
+            e.equals(cf.Data.full(shape=(2, 4), fill_value=2, dtype=int))
+        )
+
+        self.assertEqual(d[d.argmin(unravel=True)].array, 0)
+
+        d = cf.Data([4, 0, 2, 3, 0])
+        self.assertEqual(d.argmin().array, 1)
+
+        # Bad axis
+        with self.assertRaises(ValueError):
+            d.argmin(axis=d.ndim)
+
     def test_Data_percentile_median(self):
         """Test the `percentile` and `median` Data methods."""
         # ranks: a sequence of percentile rank inputs. NOTE: must

cf/test/test_Field.py

@@ -2572,6 +2572,21 @@ def test_Field_argmax(self):
         with self.assertRaises(ValueError):
             f.argmax(axis="foo")
 
+    def test_Field_argmin(self):
+        """Test the `argmin` Field method."""
+        f = cf.example_field(2)
+        i = f.argmin("T")
+        self.assertEqual(i.shape, f.shape[1:])
+
+        i = f.argmin(unravel=True)
+        self.assertIsInstance(i, tuple)
+        g = f[i]
+        self.assertEqual(g.shape, (1, 1, 1))
+
+        # Bad axis
+        with self.assertRaises(ValueError):
+            f.argmin(axis="foo")
+
     def test_Field_subspace(self):
         f = self.f
 

docs/source/class/cf.Data.rst

@@ -458,6 +458,7 @@ Searching
    :template: method.rst
 
    ~cf.Data.argmax
+   ~cf.Data.argmin
    ~cf.Data.where
 
 Counting

docs/source/class/cf.Field.rst

@@ -241,6 +241,7 @@ Data
    :template: method.rst
 
    ~cf.Field.argmax
+   ~cf.Field.argmin
    ~cf.Field.where
 
 Miscellaneous data operations