Merge pull request #10822 from george-galvin/convolution-issue

Made convolve work with units, closing issue #10811

CHANGES.rst

@@ -13,6 +13,9 @@ astropy.constants
 astropy.convolution
 ^^^^^^^^^^^^^^^^^^^
 
+- Methods ``convolve`` and ``convolve_fft`` both now return Quantity arrays if user
+  input is given in one. [#10822]
+
 astropy.coordinates
 ^^^^^^^^^^^^^^^^^^^
 

astropy/convolution/convolve.py

@@ -353,6 +353,10 @@ def convolve(array, kernel, boundary='fill', fill_value=0.,
         result[initially_nan] = np.nan
 
     # Convert result to original data type
+    array_unit = getattr(passed_array, "unit", None)
+    if array_unit is not None:
+        result <<= array_unit
+
     if isinstance(passed_array, Kernel):
         if isinstance(passed_array, Kernel1D):
             new_result = Kernel1D(array=result)
@@ -567,6 +571,9 @@ def convolve_fft(array, kernel, boundary='fill', fill_value=0.,
     if nan_treatment not in ('interpolate', 'fill'):
         raise ValueError("nan_treatment must be one of 'interpolate','fill'")
 
+    #Get array quantity if it exists
+    array_unit = getattr(array, "unit", None)
+
     # Convert array dtype to complex
     # and ensure that list inputs become arrays
     array = _copy_input_if_needed(array, dtype=complex, order='C',
@@ -759,6 +766,9 @@ def convolve_fft(array, kernel, boundary='fill', fill_value=0.,
 
     fftmult *= kernel_scale
 
+    if array_unit is not None:
+        fftmult <<= array_unit
+
     if return_fft:
         return fftmult
 

astropy/convolution/tests/test_convolve.py

@@ -43,6 +43,7 @@
 
 
 class TestConvolve1D:
+
     def test_list(self):
         """
         Test that convolve works correctly when inputs are lists
@@ -65,6 +66,26 @@ def test_tuple(self):
         assert_array_almost_equal_nulp(z,
             np.array([0., 3.6, 5., 5.6, 5.6, 6.8, 0.]), 10)
 
+    @pytest.mark.parametrize(('boundary', 'nan_treatment',
+                              'normalize_kernel', 'preserve_nan', 'dtype'),
+                             itertools.product(BOUNDARY_OPTIONS,
+                                               NANHANDLING_OPTIONS,
+                                               NORMALIZE_OPTIONS,
+                                               PRESERVE_NAN_OPTIONS,
+                                               VALID_DTYPES))
+    def test_quantity(self, boundary, nan_treatment,
+                      normalize_kernel, preserve_nan, dtype):
+        """
+        Test that convolve works correctly when input array is a Quantity
+        """
+
+        x = np.array([1, 4, 5, 6, 5, 7, 8], dtype=dtype) * u.ph
+        y = np.array([0.2, 0.6, 0.2], dtype=dtype)
+        z = convolve(x, y, boundary=boundary, nan_treatment=nan_treatment,
+                          normalize_kernel=normalize_kernel, preserve_nan=preserve_nan)
+
+        assert x.unit == z.unit
+
     @pytest.mark.parametrize(('boundary', 'nan_treatment',
                               'normalize_kernel', 'preserve_nan', 'dtype'),
                              itertools.product(BOUNDARY_OPTIONS,

astropy/convolution/tests/test_convolve_fft.py

@@ -8,6 +8,7 @@
 
 from astropy.convolution.convolve import convolve_fft, convolve
 from astropy.utils.exceptions import AstropyUserWarning
+from astropy import units as u
 from astropy.utils.compat.context import nullcontext
 
 VALID_DTYPES = ('>f4', '<f4', '>f8', '<f8')
@@ -69,6 +70,22 @@ def assert_floatclose(x, y):
 
 class TestConvolve1D:
 
+    @pytest.mark.parametrize(option_names, options)
+    def test_quantity(self, boundary, nan_treatment, normalize_kernel):
+        """
+        Test that convolve_fft works correctly when input array is a Quantity
+        """
+
+        x = np.array([1., 4., 5., 6., 5., 7., 8.], dtype='float64') * u.ph
+        y = np.array([0.2, 0.6, 0.2], dtype='float64')
+
+        with expected_boundary_warning(boundary=boundary):
+            z = convolve_fft(x, y, boundary=boundary,
+                             nan_treatment=nan_treatment,
+                             normalize_kernel=normalize_kernel)
+
+        assert x.unit == z.unit
+
     @pytest.mark.parametrize(option_names, options)
     def test_unity_1_none(self, boundary, nan_treatment, normalize_kernel):
         '''