From 4bac36b3823e693d8e421f16a36fe3846efd0e39 Mon Sep 17 00:00:00 2001
From: Oscar Esteban <code@oscaresteban.es>
Date: Fri, 16 Dec 2022 17:06:08 +0100
Subject: [PATCH] enh: add implementation of smoothing and decimating

---
 src/eddymotion/data/filtering.py | 87 ++++++++++++++++++++++++++++++++
 src/eddymotion/estimator.py      |  3 +-
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/src/eddymotion/data/filtering.py b/src/eddymotion/data/filtering.py
index c84cc02b..345e527f 100644
--- a/src/eddymotion/data/filtering.py
+++ b/src/eddymotion/data/filtering.py
@@ -23,6 +23,93 @@
 """Filtering data."""
 
 
+def gaussian_filter(data, vox_width):
+    """
+    Apply a Gaussian smoothing filter of a given width (in voxels)
+
+    Parameters
+    ----------
+    data : :obj:`numpy.ndarray`
+        The input image's data array
+    vox_width : :obj:`numbers.Number` or :obj:`tuple` or :obj:`list`
+        The smoothing kernel width in voxels
+
+    Returns
+    -------
+    data : :obj:`numpy.ndarray`
+        The smoothed dataset
+
+    """
+    from numbers import Number
+    import numpy as np
+    from scipy.ndimage import gaussian_filter as _gs
+
+    data = np.squeeze(data)  # drop unused dimensions
+    ndim = data.ndim
+
+    if isinstance(vox_width, Number):
+        vox_width = tuple([vox_width] * min(3, ndim))
+
+    # Do not smooth across time/orientation
+    if ndim == 4 and len(vox_width) == 3:
+        vox_width = (*vox_width, 0)
+
+    return _gs(data, vox_width)
+
+
+def decimate(in_file, factor, smooth=True, order=3, nonnegative=True):
+    from numbers import Number
+    import numpy as np
+    from scipy.ndimage import map_coordinates
+    import nibabel as nb
+
+    imnii = nb.load(in_file)
+    data = np.squeeze(imnii.get_fdata())
+    datashape = data.shape
+    ndim = data.ndim
+
+    if isinstance(factor, Number):
+        factor = tuple([factor] * min(3, ndim))
+
+    if ndim == 4 and len(factor) == 3:
+        factor = (*factor, 0)
+
+    if smooth:
+        if smooth is True:
+            smooth = factor
+
+        data = gaussian_filter(data, smooth)
+
+    down_grid = np.array(
+        np.meshgrid(
+            *[np.arange(_s, step=int(_f) or 1) for _s, _f in zip(datashape, factor)],
+            indexing="ij",
+        )
+    )
+    new_shape = down_grid.shape[1:]
+    newaffine = imnii.affine.copy()
+    newaffine[:3, :3] = np.array(factor[:3]) * newaffine[:3, :3]
+    # newaffine[:3, 3] += imnii.affine[:3, :3] @ (0.5 / np.array(factor[:3], dtype="float32"))
+
+    # Resample data in the new grid
+    resampled = map_coordinates(
+        data,
+        down_grid.reshape((ndim, np.prod(new_shape))),
+        order=order,
+        mode="constant",
+        cval=0,
+        prefilter=True,
+    ).reshape(new_shape)
+
+    if order > 2 and nonnegative:
+        resampled[resampled < 0] = 0
+
+    newnii = nb.Nifti1Image(resampled, newaffine, imnii.header)
+    newnii.set_sform(newaffine, code=1)
+    newnii.set_qform(newaffine, code=1)
+    return newnii
+
+
 def advanced_clip(
     data, p_min=35, p_max=99.98, nonnegative=True, dtype="int16", invert=False
 ):
diff --git a/src/eddymotion/estimator.py b/src/eddymotion/estimator.py
index 4f256cc6..3fe7bdab 100644
--- a/src/eddymotion/estimator.py
+++ b/src/eddymotion/estimator.py
@@ -270,12 +270,13 @@ def _prepare_kwargs(dwdata, kwargs):
     kwargs : :obj:`dict`
         Keyword arguments.
     """
+    from eddymotion.data.filtering import advanced_clip as _advanced_clip
 
     if dwdata.brainmask is not None:
         kwargs["mask"] = dwdata.brainmask
 
     if hasattr(dwdata, "bzero") and dwdata.bzero is not None:
-        kwargs["S0"] = dwdata.bzero
+        kwargs["S0"] = _advanced_clip(dwdata.bzero)
 
     if hasattr(dwdata, "gradients"):
         kwargs["gtab"] = dwdata.gradients