[FIX] Headers not carried over by math_img

nilearn/conftest.py

@@ -346,3 +346,55 @@ def img_atlas(shape_3d_default, affine_mni):
             "csf": 3,
         },
     }
+
+
+# ------------------ 4D IMAGES, EDIT HEADER -----------------#
+
+
+@pytest.fixture
+def img_4d_ones_eye_default_header():
+    """Return a 4D Nifti1Image with default header.
+
+    The header is created by new_img_like and is not modified. The image is
+    filled with ones and has an identity affine.
+    """
+    img = image.new_img_like(
+        _img_ones(_shape_4d_default(), _affine_eye()),
+        data=_img_ones(_shape_4d_default(), _affine_eye()).get_fdata(),
+        copy_header=False,
+    )
+    return img
+
+
+@pytest.fixture
+def img_4d_ones_eye_tr2():
+    """Return a 4D Nifti1Image with otherwise default header, except TR 2.0.
+
+    The header is the default one created by new_img_like, but the TR is
+    changed to 2.0. The image is filled with ones and has an identity affine.
+    """
+    img = image.new_img_like(
+        _img_ones(_shape_4d_default(), _affine_eye()),
+        data=_img_ones(_shape_4d_default(), _affine_eye()).get_fdata(),
+        copy_header=False,
+    )
+    # Change the TR
+    header = img.header.copy()
+    header["pixdim"][4] = 2.0
+    return Nifti1Image(img.get_fdata(), img.affine, header=header)
+
+
+@pytest.fixture
+def img_4d_mni_tr2():
+    """Return a 4D Nifti1Image with MNI affine and header, and TR 2.0.
+
+    The header has the MNI affine, and the TR is changed to 2.0. The image is
+    filled with random numbers.
+    """
+    img = image.new_img_like(
+        _img_4d_mni(), data=_img_4d_mni().get_fdata(), copy_header=True
+    )
+    # Change the TR
+    header = img.header.copy()
+    header["pixdim"][4] = 2.0
+    return Nifti1Image(img.get_fdata(), img.affine, header=header)

nilearn/image/image.py

@@ -989,7 +989,7 @@ def threshold_img(
     return thresholded_img
 
 
-def math_img(formula, **imgs):
+def math_img(formula, copy_header_from=None, **imgs):
     """Interpret a numpy based string formula using niimg in named parameters.
 
     .. versionadded:: 0.2.3
@@ -1000,6 +1000,15 @@ def math_img(formula, **imgs):
         The mathematical formula to apply to image internal data. It can use
         numpy imported as 'np'.
 
+    copy_header_from : :obj:`str`, default=None
+        Takes the variable name of one of the images in the formula.
+        The header of this image will be copied to the result of the formula.
+        Note that the result image and the image to copy the header from,
+        should have the same number of dimensions. If None, the default
+        :class:`~nibabel.nifti1.Nifti1Header` is used.
+
+        .. versionadded:: 0.10.4.dev
+
     imgs : images (:class:`~nibabel.nifti1.Nifti1Image` or file names)
         Keyword arguments corresponding to the variables in the formula as
         Nifti images. All input images should have the same geometry (shape,
@@ -1033,6 +1042,30 @@ def math_img(formula, **imgs):
      >>> result_img = math_img("img1 + img2",
      ...                       img1=anatomical_image, img2=log_img)
 
+    The result image will have the same shape and affine as the input images;
+    but might have different header information, specifically the TR value,
+    see :gh:`2645`.
+
+    .. versionadded:: 0.10.4.dev
+
+    We can now copy the header from one of the input images::
+
+     >>> from nilearn.image import load_img
+     >>> haxby_data = datasets.fetch_haxby()
+     >>> bold_img = haxby_data.func[0]
+     >>> # check input image TR
+     >>> print(load_img(bold_img).header.get_zooms()[3])
+         2.5
+     >>> result_img_without_header = math_img("img**2", img=bold_img)
+     >>> # check result image TR
+     >>> print(result_img_without_header.header.get_zooms()[3])
+         1.0
+     >>> result_img_with_header = math_img("img**2", img=bold_img,
+     ...                                   copy_header_from="img")
+     >>> # result image TR is now same as input image TR
+     >>> print(result_img_with_header.header.get_zooms()[3])
+         2.5
+
     Notes
     -----
     This function is the Python equivalent of ImCal in SPM or fslmaths
@@ -1068,7 +1101,20 @@ def math_img(formula, **imgs):
         ) + exc.args
         raise
 
-    return new_img_like(niimg, result, niimg.affine)
+    # check whether to copy header from one of the input images
+    if copy_header_from is not None:
+        niimg = check_niimg(imgs[copy_header_from])
+        # only copy the header if the result and the input image to copy the
+        # header from have the same shape
+        if result.ndim != niimg.ndim:
+            raise ValueError(
+                "Cannot copy the header. "
+                "The result of the formula has a different number of "
+                "dimensions than the image to copy the header from."
+            )
+        return new_img_like(niimg, result, niimg.affine, copy_header=True)
+    else:
+        return new_img_like(niimg, result, niimg.affine)
 
 
 def binarize_img(img, threshold=0, mask_img=None, two_sided=True):

nilearn/image/tests/test_image.py

@@ -818,6 +818,18 @@ def test_math_img_exceptions(affine_eye, img_4d_ones_eye):
     ):
         math_img(bad_formula, img1=img1, img3=img3)
 
+    # Test copy_header_from parameter
+    # Copying header from 4d image to a result that is 3d should raise a
+    # ValueError
+    formula = "np.mean(img1, axis=-1) - np.mean(img3, axis=-1)"
+    with pytest.raises(ValueError, match="Cannot copy the header."):
+        math_img(formula, img1=img1, img3=img3, copy_header_from="img1")
+
+    # Passing an 'img*' variable (to copy_header_from) that is not in the
+    # formula or an img* argument should raise a KeyError exception.
+    with pytest.raises(KeyError):
+        math_img(formula, img1=img1, img3=img3, copy_header_from="img2")
+
 
 def test_math_img(
     affine_eye, img_4d_ones_eye, img_4d_zeros_eye, shape_3d_default, tmp_path
@@ -837,6 +849,55 @@ def test_math_img(
         assert result.shape == expected_result.shape
 
 
+def test_math_img_copied_header(
+    img_4d_ones_eye_default_header, img_4d_ones_eye_tr2, img_4d_mni_tr2
+):
+    img_eye_default = img_4d_ones_eye_default_header
+    img_eye_tr2 = img_4d_ones_eye_tr2
+    img_mni_tr2 = img_4d_mni_tr2
+
+    # case where data values are not changed and header values are not copied
+    # the result should have default header values
+    formula_no_change = "img * 1"
+    # using img_eye_tr2 in the formula
+    result = math_img(
+        formula_no_change, img=img_eye_tr2, copy_header_from=None
+    )
+    # header values should instead be same as for img_eye_default that has
+    # default header values
+    assert result.header.__eq__(img_eye_default.header)
+
+    # case where data values are not changed but header values are copied
+    # the result should have the same header values as img_mni_tr2
+    result = math_img(
+        formula_no_change, img=img_mni_tr2, copy_header_from="img"
+    )
+    # all header values should be the same
+    assert result.header.__eq__(img_mni_tr2.header)
+
+    # case where data values are changed and header values are copied from one
+    # of the input images
+    # the result should have the same header values as img_eye_tr2, except for
+    # cal_max and cal_min that should be different
+    formula_change_min_max = "img1 - img2"
+    result = math_img(
+        formula_change_min_max,
+        img1=img_eye_default,
+        img2=img_eye_tr2,
+        copy_header_from="img2",
+    )
+    for key in img_mni_tr2.header.keys():
+        # cal_max and cal_min should be different in result
+        if key in ["cal_max", "cal_min"]:
+            assert result.header[key] != img_mni_tr2.header[key]
+        # other header values should be the same
+        else:
+            if isinstance(result.header[key], np.ndarray):
+                assert_array_equal(result.header[key], img_eye_tr2.header[key])
+            else:
+                assert result.header[key] == img_eye_tr2.header[key]
+
+
 def test_binarize_img(img_4d_rand_eye):
     # Test that all output values are 1.
     img1 = binarize_img(img_4d_rand_eye)