Fix bug that caused Data._axes to be incorrect after a call to cf.Field.collapse

Author: davidhassell

Changelog.rst

@@ -13,6 +13,9 @@ version 3.17.0
   ``key`` (https://github.com/NCAS-CMS/cf-python/issues/802)
 * New keyword parameter to `cf.histogram`: ``density``
   (https://github.com/NCAS-CMS/cf-python/issues/794)
+* Fix bug that caused `Data._axes` to be incorrect after a call to
+  `cf.Field.collapse`
+  (https://github.com/NCAS-CMS/cf-python/issues/857)
 * Changed dependency: ``Python>=3.9.0``
 * Changed dependency: ``numpy>=2.0.0``
 * Changed dependency: ``cfdm>=1.12.0.0, <1.12.1.0``

cf/data/utils.py

@@ -402,8 +402,14 @@ def collapse(
             weights, axis)``.
 
     """
+    original_size = d.size
+    if axis is None:
+        axis = range(d.ndim)
+    else:
+        axis = d._parse_axes(axis)
+
     kwargs = {
-        "axis": axis,
+        "axis": tuple(axis),
         "keepdims": keepdims,
         "split_every": split_every,
         "mtol": mtol,
@@ -424,6 +430,14 @@ def collapse(
     dx = func(dx, **kwargs)
     d._set_dask(dx)
 
+    if not keepdims:
+        # Remove collapsed axis names
+        d._axes = [a for i, a in enumerate(d._axes) if i not in axis]
+
+    if d.size != original_size:
+        # Remove the out-dated HDF5 chunking strategy
+        d.nc_clear_hdf5_chunksizes()
+
     return d, weights
 
 

cf/test/test_Data.py

@@ -4657,6 +4657,33 @@ def test_Data_is_masked(self):
         self.assertTrue(d[0].is_masked)
         self.assertFalse(d[1].is_masked)
 
+    def test_Data_collapse_axes_hdf_chunks(self):
+        """Test that _axes and hdf_chunks are updated after a collapse."""
+        d = cf.Data([[1, 2, 3, 4]])
+        chunks = d.shape
+        d.nc_set_hdf5_chunksizes(chunks)
+        e = d.mean(axes=1)
+        self.assertEqual(d._axes, ("dim0", "dim1"))
+        self.assertEqual(d.nc_hdf5_chunksizes(), chunks)
+
+        e = d.mean(axes=1)
+        self.assertNotEqual(e.size, d.size)
+        self.assertEqual(e._axes, d._axes)
+        self.assertEqual(e.nc_hdf5_chunksizes(), None)
+
+        e = d.mean(axes=1, squeeze=True)
+        self.assertEqual(e._axes, d._axes[:1])
+        self.assertEqual(e.nc_hdf5_chunksizes(), None)
+
+        e = d.mean(axes=0)
+        self.assertEqual(e.size, d.size)
+        self.assertEqual(e._axes, d._axes)
+        self.assertEqual(e.nc_hdf5_chunksizes(), chunks)
+
+        e = d.mean(axes=0, squeeze=True)
+        self.assertEqual(e._axes, d._axes[1:])
+        self.assertEqual(e.nc_hdf5_chunksizes(), chunks)
+
 
 if __name__ == "__main__":
     print("Run date:", datetime.datetime.now())

docs/source/recipes/plot_08_recipe.py

@@ -9,11 +9,10 @@
 # 1. Import cf-python, cf-plot, numpy and scipy.stats:
 
 import cfplot as cfp
-import cf
-
 import numpy as np
 import scipy.stats as stats
 
+import cf
 
 # %%
 # 2. Three functions are defined:

docs/source/recipes/plot_12_recipe.py

@@ -13,8 +13,8 @@
 # %%
 # 1. Import cf-python, cf-plot and matplotlib.pyplot:
 
-import matplotlib.pyplot as plt
 import cfplot as cfp
+import matplotlib.pyplot as plt
 
 import cf
 

docs/source/recipes/plot_13_recipe.py

@@ -18,13 +18,11 @@
 # in next steps.
 
 import cartopy.crs as ccrs
-import matplotlib.patches as mpatches
-
 import cfplot as cfp
+import matplotlib.patches as mpatches
 
 import cf
 
-
 # %%
 # 2. Read and select the SST by index and look at its contents:
 sst = cf.read("~/recipes/ERA5_monthly_averaged_SST.nc")[0]

docs/source/recipes/plot_17_recipe.py

@@ -11,8 +11,8 @@
 # %%
 # 1. Import cf-python and cf-plot:
 
-import matplotlib.pyplot as plt
 import cfplot as cfp
+import matplotlib.pyplot as plt
 
 import cf
 

docs/source/recipes/plot_18_recipe.py

@@ -10,15 +10,15 @@
 
 """
 
+import cfplot as cfp
+
 # %%
 # 1. Import cf-python, cf-plot and other required packages:
 import matplotlib.pyplot as plt
 import scipy.stats.mstats as mstats
-import cfplot as cfp
 
 import cf
 
-
 # %%
 # 2. Read the data in and unpack the Fields from FieldLists using indexing.
 # In our example We are investigating the influence of the land height on

docs/source/recipes/plot_19_recipe.py

@@ -9,10 +9,11 @@
 maxima.
 """
 
+import cfplot as cfp
+
 # %%
 # 1. Import cf-python, cf-plot and other required packages:
 import matplotlib.pyplot as plt
-import cfplot as cfp
 
 import cf