Repository: nilearn/nilearn

Issue:
[BUG] clean_img with sample_mask passed to signal.clean fails because output dimensions are not the same as the input
### Is there an existing issue for this?

- [X] I have searched the existing issues

### Operating system

- [X] Linux
- [ ] Mac
- [ ] Windows

### Operating system version

Red Hat Enterprise Linux 8.8 (Ootpa)

### Python version

- [ ] 3.11
- [X] 3.10
- [ ] 3.9
- [ ] 3.8
- [ ] 3.7

### nilearn version

0.10.1

### Expected behavior

I'm trying to apply a nuisance regression to data preprocessed with fmriprep. I am reading in certain confounds and, according to threshold criteria, creating a `sample_mask` with `nilearn.interfaces.fmriprep.load_confounds`. 

I would like to clean the data according to the algorithm presented in `nilearn.signal.clean`, but I want to input the nifti, so I use `nilearn.image.clean_image`. Since `image.clean_image` passes data to `signal.clean`, and `signal.clean` scrubs the data if you give it `sample_mask`, I expect that giving `image.clean_img` a `sample_mask` as a kwarg should result in cleaned and scrubbed data. 

So the output should be fewer timepoints than the input. This is possible in `signal.clean` but the current help of `image.clean_img` says that the output is the same dimensions as the input.

### Current behavior & error messages

This is what I got:

```bash
nifti.shape
     (97, 115, 97, 391)

len(sample_mask)
     310

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_2542916/2394484833.py in <module>
      8 nifti = nib.load(f'{path}/{file}')
      9 confounds_out, sample_mask = load_confounds(f'{path}/{file}', strategy=('motion','wm_csf', 'scrub'), motion='full', scrub=5,fd_threshold=0.3, std_dvars_threshold=3, wm_csf='basic', compcor='anat_combined', n_compcor=5, demean=True)
---> 10 clean_nifti = image.clean_img(nifti, detrend=True, standardize='zscore_sample', confounds=confounds_out, standardize_confounds=True, filter='butterworth', low_pass=0.1, high_pass=0.01, t_r=0.8, **{'clean__sample_mask':sample_mask})

~/miniconda3/lib/python3.7/site-packages/nilearn/_utils/helpers.py in wrapper(*args, **kwargs)
     40                                                      kwargs
     41                                                      )
---> 42             return func(*args, **kwargs)
     43 
     44         return wrapper

~/miniconda3/lib/python3.7/site-packages/nilearn/image/image.py in clean_img(imgs, runs, detrend, standardize, confounds, low_pass, high_pass, t_r, ensure_finite, mask_img, **kwargs)
   1287     else:
   1288         imgs_ = new_img_like(
-> 1289             imgs_, data.T.reshape(imgs_.shape), copy_header=True
   1290         )
   1291 

ValueError: cannot reshape array of size 335430850 into shape (97,115,97,391)
```


### Steps and code to reproduce bug

```python
from nilearn.interfaces.fmriprep import load_confounds
from nilearn import image
import nibabel as nib

# image should have > 0 timepoints to be scrubbed
file ='preproc_bold.nii.gz'
path = 'path/to/func/'

nifti = nib.load(f'{path}/{file}')
confounds_out, sample_mask = load_confounds(f'{path}/{file}', strategy=('motion','wm_csf', 'scrub'), motion='full', scrub=5,fd_threshold=0.3, std_dvars_threshold=3, wm_csf='basic', compcor='anat_combined', n_compcor=5, demean=True)
clean_nifti = image.clean_img(nifti, detrend=True, standardize='zscore_sample', confounds=confounds_out, standardize_confounds=True, filter='butterworth', low_pass=0.1, high_pass=0.01, t_r=0.8, **{'clean__sample_mask':sample_mask})
```



Discussion:

Resolving patch:
diff --git a/nilearn/image/image.py b/nilearn/image/image.py
--- a/nilearn/image/image.py
+++ b/nilearn/image/image.py
@@ -1280,6 +1280,11 @@ def clean_img(
     # Put results back into Niimg-like object
     if mask_img is not None:
         imgs_ = masking.unmask(data, mask_img)
+    elif "sample_mask" in clean_kwargs:
+        sample_shape = imgs_.shape[:3] + clean_kwargs["sample_mask"].shape
+        imgs_ = new_img_like(
+            imgs_, data.T.reshape(sample_shape), copy_header=True
+        )
     else:
         imgs_ = new_img_like(
             imgs_, data.T.reshape(imgs_.shape), copy_header=True


Test patch:
diff --git a/nilearn/image/tests/test_image.py b/nilearn/image/tests/test_image.py
--- a/nilearn/image/tests/test_image.py
+++ b/nilearn/image/tests/test_image.py
@@ -20,6 +20,7 @@
 from nilearn import signal
 from nilearn._utils import niimg_conversions, testing
 from nilearn._utils.data_gen import (
+    _basic_confounds,
     generate_fake_fmri,
     generate_labeled_regions,
     generate_maps,
@@ -989,3 +990,41 @@ def test_new_img_like_boolean_data(affine_eye, image, shape_3d_default, rng):
     out_img = new_img_like(in_img, data=in_img.get_fdata() > 0.5)
 
     assert get_data(out_img).dtype == "uint8"
+
+
+def test_clean_img_sample_mask(img_4d_rand_eye):
+    """Check sample mask can be passed as a kwarg and used correctly."""
+    length = 10
+    confounds = _basic_confounds(length)
+    # exclude last time point
+    sample_mask = np.arange(length - 1)
+
+    img = image.clean_img(
+        img_4d_rand_eye,
+        confounds=confounds,
+        **{"clean__sample_mask": sample_mask},
+    )
+    # original shape is (10, 10, 10, 10)
+    assert img.shape == (10, 10, 10, 9)
+
+
+def test_clean_img_sample_mask_mask_img(shape_3d_default):
+    """Check sample_mask and mask_img can be correctly used together."""
+    length = 10
+    confounds = _basic_confounds(length)
+    img_4d, mask_img = generate_fake_fmri(
+        shape=shape_3d_default, length=length
+    )
+
+    # exclude last time point
+    sample_mask = np.arange(length - 1)
+
+    # test with sample mask
+    img = image.clean_img(
+        img_4d,
+        confounds=confounds,
+        mask_img=mask_img,
+        **{"clean__sample_mask": sample_mask},
+    )
+    # original shape is (10, 10, 10, 10)
+    assert img.shape == (10, 10, 10, 9)