allow general filter use intead of just median filters

spectral_cube/dask_spectral_cube.py

@@ -902,13 +902,22 @@ def spectral_smooth(array):
 
     @add_save_to_tmp_dir_option
     def spectral_smooth_median(self, ksize, raise_error_jybm=True, **kwargs):
+        return self.spectral_filter(ksize,
+                function=ndimage.filters.median_filter, raise_error_jybm=True,
+                **kwargs)
+
+    @add_save_to_tmp_dir_option
+    def spectral_filter(self, ksize, filter, raise_error_jybm=True,
+                        **kwargs):
         """
-        Smooth the cube along the spectral dimension
+        Smooth the cube along the spectral dimension using a filter
 
         Parameters
         ----------
         ksize : int
             Size of the median filter (scipy.ndimage.filters.median_filter)
+        filter : function
+            A filter from scipy.ndimage.filters
         save_to_tmp_dir : bool
             If `True`, the computation will be carried out straight away and
             saved to a temporary directory. This can improve performance,
@@ -928,7 +937,7 @@ def spectral_smooth_median(self, ksize, raise_error_jybm=True, **kwargs):
             raise TypeError('ksize should be an integer (got {0})'.format(ksize))
 
         def median_filter_wrapper(img, **kwargs):
-            return ndimage.median_filter(img, (ksize, 1, 1), **kwargs)
+            return filter(img, (ksize, 1, 1), **kwargs)
 
         return self.apply_function_parallel_spectral(median_filter_wrapper,
                                                      accepts_chunks=True)

spectral_cube/spectral_cube.py

@@ -2701,14 +2701,42 @@ def reproject(self, header, order='bilinear', use_memmap=False,
 
 
     @parallel_docstring
-    def spatial_smooth_median(self, ksize, update_function=None, raise_error_jybm=True, **kwargs):
+    def spatial_smooth_median(self, ksize, update_function=None, raise_error_jybm=True,
+                              filter=ndimage.filters.median_filter, **kwargs):
         """
         Smooth the image in each spatial-spatial plane of the cube using a median filter.
 
         Parameters
         ----------
         ksize : int
-            Size of the median filter (scipy.ndimage.filters.median_filter)
+            Size of the median filter (scipy.ndimage.filters.*_filter)
+        filter : function
+            A filter from scipy.ndimage.filters
+        update_function : method
+            Method that is called to update an external progressbar
+            If provided, it disables the default `astropy.utils.console.ProgressBar`
+        raise_error_jybm : bool, optional
+            Raises a `~spectral_cube.utils.BeamUnitsError` when smoothing a cube in Jy/beam units,
+            since the brightness is dependent on the spatial resolution.
+        kwargs : dict
+            Passed to the convolve function
+        """
+        return self.spatial_filter(ksize, filter,
+                update_function=update_funciton,
+                raise_error_jybeam=raise_error_jybeam, **kwargs)
+
+
+    @parallel_docstring
+    def spatial_filter(self, ksize, filter, update_function=None, raise_error_jybm=True, **kwargs):
+        """
+        Smooth the image in each spatial-spatial plane of the cube using a filter.
+
+        Parameters
+        ----------
+        ksize : int
+            Size of the median filter (scipy.ndimage.filters.*_filter)
+        filter : function
+            A filter from scipy.ndimage.filters
         update_function : method
             Method that is called to update an external progressbar
             If provided, it disables the default `astropy.utils.console.ProgressBar`
@@ -2724,7 +2752,7 @@ def spatial_smooth_median(self, ksize, update_function=None, raise_error_jybm=Tr
         self.check_jybeam_smoothing(raise_error_jybm=raise_error_jybm)
 
         def _msmooth_image(im, **kwargs):
-            return ndimage.filters.median_filter(im, size=ksize, **kwargs)
+            return filter(im, size=ksize, **kwargs)
 
         newcube = self.apply_function_parallel_spatial(_msmooth_image,
                                                        **kwargs)

spectral_cube/tests/test_spectral_cube.py

@@ -2467,6 +2467,23 @@ def test_spatial_smooth_median(data_adv, use_dask):
     np.testing.assert_almost_equal(cube_median[2].value, result2)
 
 
+@pytest.mark.parametrize('num_cores', (None, 1))
+def test_spectral_smooth_maxfilter(num_cores, data_adv, use_dask):
+
+    pytest.importorskip('scipy.ndimage')
+    from scipy import ndimage
+
+    cube, data = cube_and_raw(data_adv, use_dask=use_dask)
+
+    cube_spectral_max = cube.spectral_filter(3,
+            filter=ndimage.filters.maximum_filter, num_cores=num_cores)
+
+    # Check first slice
+    result = np.array([0.90388047, 0.90388047, 0.96629004, 0.96629004])
+
+    np.testing.assert_almost_equal(cube_spectral_max[:,1,1].value, result)
+
+
 @pytest.mark.parametrize('num_cores', (None, 1))
 def test_spectral_smooth_median(num_cores, data_adv, use_dask):