import os

from pathlib import Path

import tempfile

import time

from warnings import warn

import numpy as np

from numpy.lib.stride_tricks import sliding_window_view

from tqdm import tqdm

from dipy.stats.sketching import count_sketch

from dipy.testing.decorators import warning_for_keywords

from dipy.utils.logging import logger

from dipy.utils.optpkg import optional_package

sklearn, has_sklearn, _ = optional_package("sklearn")

linear_model, _, _ = optional_package("sklearn.linear_model")

def _vol_split(train, vol_idx):

"""Split the 3D volumes into the train and test set.

Parameters

----------

train : numpy.ndarray

Array of all 3D patches flattened out to be 2D.

vol_idx : int

The volume number that needs to be held out for training.

Returns

-------

cur_x : numpy.ndarray of shape (nvolumes * patch_size) x (nvoxels)

Array of patches corresponding to all volumes except the held out volume.

y : numpy.ndarray of shape (patch_size) x (nvoxels)

Array of patches corresponding to the volume that is used a target for

denoising.

"""

mask = np.zeros(train.shape[0], dtype=bool)

mask[vol_idx] = True

cur_x = train[~mask].reshape((train.shape[0] - 1) * train.shape[1], train.shape[2])

y = train[vol_idx, train.shape[1] // 2, :]

return cur_x, y

def _extract_3d_patches(arr, patch_radius):

"""Extract 3D patches from 4D DWI data.

Parameters

----------

arr : ndarray

The 4D noisy DWI data to be denoised.

patch_radius : int or array of shape (3,)

The radius of the local patch to be taken around each voxel (in

voxels).

Returns

-------

all_patches : ndarray

All 3D patches flattened out to be 2D corresponding to the each 3D

volume of the 4D DWI data.

"""

patch_radius = np.asarray(patch_radius, dtype=int)

if patch_radius.size == 1:

patch_radius = np.repeat(patch_radius, 3)

elif patch_radius.size != 3:

raise ValueError("patch_radius should have length 1 or 3")

patch_size = 2 * patch_radius + 1

dim = arr.shape[-1]

# Calculate the shape of the output array

output_shape = tuple(arr.shape[i] - 2 * patch_radius[i] for i in range(3))

total_patches = np.prod(output_shape)

patches = sliding_window_view(arr, tuple(patch_size) + (dim,))

# Reshape and transpose the patches to match the original function's output shape

all_patches = patches.reshape(total_patches, np.prod(patch_size), dim)

all_patches = all_patches.transpose(2, 1, 0)

return np.array(all_patches)

def _fit_denoising_model(train, vol_idx, model, alpha):

"""Fit a single 3D volume using a train and test phase.

Parameters

----------

train : ndarray

Array of all 3D patches flattened out to be 2D.

vol_idx : int

The volume number that needs to be held out for training.

model : str or sklearn.base.RegressorMixin

This will determine the algorithm used to solve the set of linear

equations underlying this model. If it is a string it needs to be

one of the following: {'ols', 'ridge', 'lasso'}. Otherwise,

it can be an object that inherits from

`dipy.optimize.SKLearnLinearSolver` or an object with a similar

interface from Scikit-Learn:

`sklearn.linear_model.LinearRegression`,

`sklearn.linear_model.Lasso` or `sklearn.linear_model.Ridge`

and other objects that inherit from `sklearn.base.RegressorMixin`.

alpha : float

Regularization parameter only for ridge and lasso regression models.

version : int

Version 1 or 3 of Patch2Self to use.

Returns

-------

model_instance : fitted linear model object

The fitted model instance if version is 3.

cur_x : ndarray

The patches corresponding to all volumes except the held out volume.

"""

if isinstance(model, str):

if model.lower() == "ols":

model_instance = linear_model.Ridge(copy_X=False, alpha=1e-10)

elif model.lower() == "ridge":

model_instance = linear_model.Ridge(copy_X=False, alpha=alpha)

elif model.lower() == "lasso":

model_instance = linear_model.Lasso(copy_X=False, max_iter=50, alpha=alpha)

else:

raise ValueError(

f"Invalid model string: {model}. Should be 'ols', 'ridge', or 'lasso'."

)

elif isinstance(model, linear_model.BaseEstimator):

model_instance = model

else:

raise ValueError(

"Model should either be a string or \

an instance of sklearn.linear_model BaseEstimator."

)

cur_x, y = _vol_split(train, vol_idx)

model_instance.fit(cur_x.T, y.T)

return model_instance, cur_x

def vol_denoise(

data_dict, b0_idx, dwi_idx, model, alpha, b0_denoising, verbose, tmp_dir

):

"""Denoise a single 3D volume using train and test phase.

Parameters

----------

data_dict : dict

Dictionary containing the following:

data_name : str

The name of the memmap file containing the memmaped data.

data_dtype : dtype

The dtype of the data.

data_shape : tuple

The shape of the data.

data_b0s : ndarray

Array of all 3D patches flattened out to be 2D for b0 volumes.

data_dwi : ndarray

Array of all 3D patches flattened out to be 2D for dwi volumes.

b0_idx : ndarray

The indices of the b0 volumes.

dwi_idx : ndarray

The indices of the dwi volumes.

model : sklearn.base.RegressorMixin

This is the model that is initialized from the `_fit_denoising_model` function.

alpha : float

Regularization parameter only for ridge and lasso regression models.

b0_denoising : bool

Skips denoising b0 volumes if set to False.

verbose : bool

Show progress of Patch2Self and time taken.

tmp_dir : str

The directory to save the temporary files.

Returns

-------

denoised_arr.name : str

The name of the memmap file containing the denoised array.

denoised_arr.dtype : dtype

The dtype of the denoised array.

denoised_arr.shape : tuple

The shape of the denoised array.

"""

data_shape = data_dict["data"][2]

data_tmp = np.memmap(

data_dict["data"][0],

dtype=data_dict["data"][1],

mode="r",

shape=data_dict["data"][2],

).reshape(np.prod(data_shape[:-1]), data_shape[-1])

data_b0s = data_dict["data_b0s"]

data_dwi = data_dict["data_dwi"]

p = data_tmp.shape[0] // 10

b0_counter = 0

dwi_counter = 0

start_idx = 0

denoised_arr_file = tempfile.NamedTemporaryFile(

delete=False, dir=tmp_dir, suffix="denoised_arr"

)

denoised_arr_file.close()

denoised_arr = np.memmap(

denoised_arr_file.name, dtype=data_tmp.dtype, mode="w+", shape=data_shape

)

idx_counter = 0

full_result = np.empty(

(data_shape[0], data_shape[1], data_shape[2], data_shape[3] // 5)

)

b0_idx = b0_idx

dwi_idx = dwi_idx

if data_b0s.shape[0] == 1:

b0_denoising = False

if not b0_denoising:

if verbose:

logger.info("b0 denoising skipped....")

for vol_idx in tqdm(

range(data_shape[-1]), desc="Fitting and Denoising", leave=False

):

if vol_idx in b0_idx.flatten():

if b0_denoising:

b_fit, _ = _fit_denoising_model(data_b0s, b0_counter, model, alpha)

b_matrix = np.zeros(data_tmp.shape[-1])

b_fit_coef = np.insert(b_fit.coef_, b0_counter, 0)

np.put(b_matrix, b0_idx, b_fit_coef)

result = np.zeros(data_tmp.shape[0])

for z in range(0, data_tmp.shape[0], p):

end_idx = z + p

if end_idx > z + p:

end_idx = data_tmp.shape[0]

result[z:end_idx] = (

np.matmul(np.squeeze(data_tmp[z:end_idx, :]), b_matrix)

+ b_fit.intercept_

)

full_result[..., idx_counter] = result.reshape(

data_shape[0], data_shape[1], data_shape[2]

)

idx_counter += 1

b0_counter += 1

del b_fit_coef

del b_matrix

del result

else:

full_result[..., idx_counter] = data_tmp[..., b0_counter].reshape(

data_shape[0], data_shape[1], data_shape[2]

)

b0_counter += 1

idx_counter += 1

else:

dwi_fit, _ = _fit_denoising_model(data_dwi, dwi_counter, model, alpha)

b_matrix = np.zeros(data_tmp.shape[-1])

dwi_fit_coef = np.insert(dwi_fit.coef_, dwi_counter, 0)

np.put(b_matrix, dwi_idx, dwi_fit_coef)

del dwi_fit_coef

result = np.zeros(data_tmp.shape[0])

for z in range(0, data_tmp.shape[0], p):

end_idx = z + p

if end_idx > z + p:

end_idx = data_tmp.shape[0]

result[z:end_idx] = (

np.matmul(np.squeeze(data_tmp[z:end_idx, :]), b_matrix)

+ dwi_fit.intercept_

)

full_result[..., idx_counter] = result.reshape(

data_shape[0], data_shape[1], data_shape[2]

)

idx_counter += 1

dwi_counter += 1

if idx_counter >= data_shape[-1] // 5:

denoised_arr[..., start_idx : vol_idx + 1] = full_result

start_idx = vol_idx + 1

idx_counter = 0

denoised_arr_idx = data_shape[-1] - data_shape[-1] % 5

full_result_idx = full_result.shape[-1] - data_shape[-1] % 5

denoised_arr[..., denoised_arr_idx:] = full_result[..., full_result_idx:]

del full_result

return denoised_arr_file.name, denoised_arr.dtype, denoised_arr.shape

@warning_for_keywords()

def patch2self(

data,

bvals,

*,

patch_radius=(0, 0, 0),

model="ols",

b0_threshold=50,

out_dtype=None,

alpha=1.0,

verbose=False,

b0_denoising=True,

clip_negative_vals=False,

shift_intensity=True,

tmp_dir=None,

version=3,

):

"""Patch2Self Denoiser.

See :footcite:p:`Fadnavis2020` for further details about the method.

See :footcite:p:`Fadnavis2024` for further details about the new method.

Parameters

----------

data : ndarray

The 4D noisy DWI data to be denoised.

bvals : array of shape (N,)

Array of the bvals from the DWI acquisition

patch_radius : int or array of shape (3,), optional

The radius of the local patch to be taken around each voxel (in

voxels).

model : string, or sklearn.base.RegressorMixin, optional

This will determine the algorithm used to solve the set of linear

equations underlying this model. If it is a string it needs to be

one of the following: {'ols', 'ridge', 'lasso'}. Otherwise,

it can be an object that inherits from

`dipy.optimize.SKLearnLinearSolver` or an object with a similar

interface from Scikit-Learn:

`sklearn.linear_model.LinearRegression`,

`sklearn.linear_model.Lasso` or `sklearn.linear_model.Ridge`

and other objects that inherit from `sklearn.base.RegressorMixin`.

b0_threshold : int, optional

Threshold for considering volumes as b0.

out_dtype : str or dtype, optional

The dtype for the output array. Default: output has the same dtype as

the input.

alpha : float, optional

Regularization parameter only for ridge regression model.

verbose : bool, optional

Show progress of Patch2Self and time taken.

b0_denoising : bool, optional

Skips denoising b0 volumes if set to False.

clip_negative_vals : bool, optional

Sets negative values after denoising to 0 using `np.clip`.

shift_intensity : bool, optional

Shifts the distribution of intensities per volume to give

non-negative values.

tmp_dir : str, optional

The directory to save the temporary files. If None, the temporary

files are saved in the system's default temporary directory.

version : int, optional

Version 1 or 3 of Patch2Self to use.

Returns

-------

denoised array : ndarray

This is the denoised array of the same size as that of the input data,

clipped to non-negative values.

References

----------

.. footbibliography::

"""

out_dtype, tmp_dir, patch_radius = _validate_inputs(

data, out_dtype, patch_radius, version, tmp_dir

)

if version == 1:

return _patch2self_version1(

data,

bvals,

patch_radius,

model,

b0_threshold,

out_dtype,

alpha,

verbose,

b0_denoising,

clip_negative_vals,

shift_intensity,

)

return _patch2self_version3(

data,

bvals,

model,

b0_threshold,

out_dtype,

alpha,

verbose,

b0_denoising,

clip_negative_vals,

shift_intensity,

tmp_dir,

)

def _validate_inputs(data, out_dtype, patch_radius, version, tmp_dir):

"""Validate inputs for patch2self function.

Parameters

----------

data : ndarray

The 4D noisy DWI data to be denoised.

out_dtype : str or dtype

The dtype for the output array.

patch_radius : int or array of shape (3,)

The radius of the local patch to be taken around each voxel (in

voxels).

version : int

Version 1 or 3 of Patch2Self to use.

tmp_dir : str

The directory to save the temporary files. If None, the temporary

files are saved in the system's default temporary directory.

Raises

------

ValueError

If temporary directory is not None for Patch2Self version 1.

If the patch_radius is not 0 for Patch2Self version 3.

If the temporary directory does not exist.

If the input data is not a 4D array.

Warns

-----

If the input data has less than 10 3D volumes.

Returns

-------

out_dtype : str or dtype

The dtype for the output array.

tmp_dir : str

The directory to save the temporary files. If None, the temporary

files are saved in the system's default temporary directory.

"""

if out_dtype is None:

out_dtype = data.dtype

if tmp_dir is None and version == 3:

tmp_dir = tempfile.gettempdir()

if version not in [1, 3]:

raise ValueError("Invalid version. Should be 1 or 3.")

if version == 1 and tmp_dir is not None:

raise ValueError(

"Temporary directory is not supported for Patch2Self version 1. \

Please set tmp_dir to None."

)

if patch_radius != (0, 0, 0) and version == 3:

raise ValueError(

"Patch radius is not supported for Patch2Self version 3. \

Please do not set patch_radius."

)

if isinstance(patch_radius, list) and len(patch_radius) == 1:

patch_radius = (patch_radius[0], patch_radius[0], patch_radius[0])

if isinstance(patch_radius, int):

patch_radius = (patch_radius, patch_radius, patch_radius)

if version == 3 and tmp_dir is not None and not Path(tmp_dir).exists():

raise ValueError("The temporary directory does not exist.")

if data.ndim != 4:

raise ValueError("Patch2Self can only denoise on 4D arrays.", data.shape)

if data.shape[3] < 10:

warn(

"The input data has less than 10 3D volumes. \

Patch2Self may not give optimal denoising performance.",

stacklevel=2,

)

return out_dtype, tmp_dir, patch_radius

def _patch2self_version1(

data,

bvals,

patch_radius,

model,

b0_threshold,

out_dtype,

alpha,

verbose,

b0_denoising,

clip_negative_vals,

shift_intensity,

):

"""Patch2Self Denoiser.

Parameters

----------

data : ndarray

The 4D noisy DWI data to be denoised.

bvals : array of shape (N,)

Array of the bvals from the DWI acquisition.

patch_radius : int or array of shape (3,)

The radius of the local patch to be taken around each voxel (in

voxels).

model : string, or sklearn.base.RegressorMixin

This will determine the algorithm used to solve the set of linear

equations underlying this model. If it is a string it needs to be

one of the following: {'ols', 'ridge', 'lasso'}. Otherwise,

it can be an object that inherits from

`dipy.optimize.SKLearnLinearSolver` or an object with a similar

interface from Scikit-Learn:

`sklearn.linear_model.LinearRegression`,

`sklearn.linear_model.Lasso` or `sklearn.linear_model.Ridge`

and other objects that inherit from `sklearn.base.RegressorMixin`.

b0_threshold : int

Threshold for considering volumes as b0.

out_dtype : str or dtype

The dtype for the output array. Default: output has the same dtype as

the input.

alpha : float

Regularization parameter only for ridge regression model.

verbose : bool

Show progress of Patch2Self and time taken.

b0_denoising : bool

Skips denoising b0 volumes if set to False.

clip_negative_vals : bool

Sets negative values after denoising to 0 using `np.clip`.

shift_intensity : bool

Shifts the distribution of intensities per volume to give

non-negative values.

Returns

-------

denoised array : ndarray

This is the denoised array of the same size as that of the input data,

clipped to non-negative values.

"""

# We retain float64 precision, iff the input is in this precision:

if data.dtype == np.float64:

calc_dtype = np.float64

# Otherwise, we'll calculate things in float32 (saving memory)

else:

calc_dtype = np.float32

original_shape = data.shape

if 1 in data.shape and data.shape[-1] != 1:

position = data.shape.index(1)

data = np.concatenate((data, data, data), position)

# Segregates volumes by b0 threshold

b0_idx = np.argwhere(bvals <= b0_threshold)

dwi_idx = np.argwhere(bvals > b0_threshold)

data_b0s = np.squeeze(np.take(data, b0_idx, axis=3))

data_dwi = np.squeeze(np.take(data, dwi_idx, axis=3))

# create empty arrays

denoised_b0s = np.empty(data_b0s.shape, dtype=calc_dtype)

denoised_dwi = np.empty(data_dwi.shape, dtype=calc_dtype)

denoised_arr = np.empty(data.shape, dtype=calc_dtype)

if verbose is True:

t1 = time.time()

# if only 1 b0 volume, skip denoising it

if data_b0s.ndim == 3 or not b0_denoising:

if verbose:

logger.info("b0 denoising skipped...")

denoised_b0s = data_b0s

else:

train_b0 = _extract_3d_patches(

np.pad(

data_b0s,

(

(patch_radius[0], patch_radius[0]),

(patch_radius[1], patch_radius[1]),

(patch_radius[2], patch_radius[2]),

(0, 0),

),

mode="constant",

),

patch_radius=patch_radius,

)

for vol_idx in range(0, data_b0s.shape[3]):

b0_model, cur_x = _fit_denoising_model(

train_b0, vol_idx, model, alpha=alpha

)

denoised_b0s[..., vol_idx] = b0_model.predict(cur_x.T).reshape(

data_b0s.shape[0], data_b0s.shape[1], data_b0s.shape[2]

)

if verbose is True:

logger.info(f"Denoised b0 Volume: {vol_idx}")

# Separate denoising for DWI volumes

train_dwi = _extract_3d_patches(

np.pad(

data_dwi,

(

(patch_radius[0], patch_radius[0]),

(patch_radius[1], patch_radius[1]),

(patch_radius[2], patch_radius[2]),

(0, 0),

),

mode="constant",

),

patch_radius=patch_radius,

)

# Insert the separately denoised arrays into the respective empty arrays

for vol_idx in range(0, data_dwi.shape[3]):

dwi_model, cur_x = _fit_denoising_model(train_dwi, vol_idx, model, alpha=alpha)

denoised_dwi[..., vol_idx] = dwi_model.predict(cur_x.T).reshape(

data_dwi.shape[0], data_dwi.shape[1], data_dwi.shape[2]

)

if verbose is True:

logger.info(f"Denoised DWI Volume: {vol_idx}")

if verbose is True:

t2 = time.time()

logger.info(f"Total time taken for Patch2Self: {t2 - t1:.3f} seconds")

if data_b0s.ndim == 3:

denoised_arr[:, :, :, b0_idx[0][0]] = denoised_b0s

else:

for i, idx in enumerate(b0_idx):

denoised_arr[:, :, :, idx[0]] = np.squeeze(denoised_b0s[..., i])

for i, idx in enumerate(dwi_idx):

denoised_arr[:, :, :, idx[0]] = np.squeeze(denoised_dwi[..., i])

if 1 in original_shape and original_shape[-1] != 1:

denoised_arr = np.take(denoised_arr, [0], axis=position)

denoised_arr = _apply_post_processing(

denoised_arr, shift_intensity, clip_negative_vals

)

return np.array(denoised_arr, dtype=out_dtype)

def _patch2self_version3(

data,

bvals,

model,

b0_threshold,

out_dtype,

alpha,

verbose,

b0_denoising,

clip_negative_vals,

shift_intensity,

tmp_dir,

):

"""Patch2Self Denoiser.

Parameters

----------

data : ndarray

The 4D noisy DWI data to be denoised.

bvals : array of shape (N,)

Array of the bvals from the DWI acquisition.

model : string, or sklearn.base.RegressorMixin

This will determine the algorithm used to solve the set of linear

equations underlying this model. If it is a string it needs to be

one of the following: {'ols', 'ridge', 'lasso'}. Otherwise,

it can be an object that inherits from

`dipy.optimize.SKLearnLinearSolver` or an object with a similar

interface from Scikit-Learn:

`sklearn.linear_model.LinearRegression`,

`sklearn.linear_model.Lasso` or `sklearn.linear_model.Ridge`

and other objects that inherit from `sklearn.base.RegressorMixin`.

b0_threshold : int

Threshold for considering volumes as b0.

out_dtype : str or dtype

The dtype for the output array. Default: output has the same dtype as

the input.

alpha : float

Regularization parameter only for ridge regression model.

verbose : bool

Show progress of Patch2Self and time taken.

b0_denoising : bool

Skips denoising b0 volumes if set to False.

clip_negative_vals : bool

Sets negative values after denoising to 0 using `np.clip`.

shift_intensity : bool

Shifts the distribution of intensities per volume to give

non-negative values.

tmp_dir : str

The directory to save the temporary files. If None, the temporary

files are saved in the system's default temporary directory.

Returns

-------

denoised array : ndarray

This is the denoised array of the same size as that of the input data,

clipped to non-negative values.

"""

if data.dtype == np.float64:

calc_dtype = np.float64

else:

calc_dtype = np.float32

tmp_file = tempfile.NamedTemporaryFile(delete=False, dir=tmp_dir, suffix="tmp_file")

tmp_file.close()

tmp = np.memmap(

tmp_file.name,

dtype=calc_dtype,

mode="w+",

shape=(data.shape[0], data.shape[1], data.shape[2], data.shape[3]),

)

p = data.shape[-1] // 5

idx_start = 0

for z in range(0, data.shape[3], p):

end_idx = z + p

if end_idx > data.shape[3]:

end_idx = data.shape[3]

if verbose:

logger.info(f"Loading data from {idx_start} to {end_idx}")

tmp[..., idx_start:end_idx] = data[..., idx_start:end_idx]

idx_start = end_idx

sketch_rows = int(0.30 * data.shape[0] * data.shape[1] * data.shape[2])

sketched_matrix_name, sketched_matrix_dtype, sketched_matrix_shape = count_sketch(

tmp_file.name,

calc_dtype,

tmp.shape,

sketch_rows=sketch_rows,

tmp_dir=tmp_dir,

)

sketched_matrix = np.memmap(

sketched_matrix_name,

dtype=sketched_matrix_dtype,

mode="r",

shape=sketched_matrix_shape,

).T

if verbose:

logger.info("Sketching done.")

b0_idx = np.argwhere(bvals <= b0_threshold)

dwi_idx = np.argwhere(bvals > b0_threshold)

data_b0s = np.take(np.squeeze(sketched_matrix), b0_idx, axis=0)

data_dwi = np.take(np.squeeze(sketched_matrix), dwi_idx, axis=0)

data_dict = {

"data": [tmp_file.name, calc_dtype, tmp.shape],

"data_b0s": data_b0s,

"data_dwi": data_dwi,

}

if verbose:

t1 = time.time()

del sketched_matrix

os.unlink(sketched_matrix_name)

denoised_arr_name, calc_dtype, denoised_arr_shape = vol_denoise(

data_dict,

b0_idx,

dwi_idx,

model,

alpha,

b0_denoising,

verbose,

tmp_dir=tmp_dir,

)

denoised_arr = np.memmap(

denoised_arr_name,

dtype=calc_dtype,

mode="r+",

shape=denoised_arr_shape,

)

if verbose:

t2 = time.time()

logger.info(f"Time taken for Patch2Self: {t2 - t1:.3f} seconds.")

denoised_arr = _apply_post_processing(

denoised_arr, shift_intensity, clip_negative_vals

)

del tmp

os.unlink(data_dict["data"][0])

result = np.array(denoised_arr, dtype=out_dtype)

del denoised_arr

os.unlink(denoised_arr_name)

return result

def _apply_post_processing(denoised_arr, shift_intensity, clip_negative_vals):

"""Apply post-processing steps such as clipping and shifting intensities.

Parameters

----------

denoised_arr : ndarray

The denoised array.

shift_intensity : bool

Shifts the distribution of intensities per volume to give

non-negative values.

clip_negative_vals : bool

Sets negative values after denoising to 0 using `np.clip`.

Returns

-------

denoised_arr : ndarray

The denoised array with post-processing applied.

"""

if shift_intensity and not clip_negative_vals:

for i in range(denoised_arr.shape[-1]):

shift = np.min(denoised_arr[..., i]) - np.min(denoised_arr[..., i])

denoised_arr[..., i] += shift

elif clip_negative_vals and not shift_intensity:

denoised_arr.clip(min=0, out=denoised_arr)

elif clip_negative_vals and shift_intensity:

warn(

"Both `clip_negative_vals` and `shift_intensity` cannot be True. \

Defaulting to `clip_negative_vals`...",

stacklevel=2,

)

denoised_arr.clip(min=0, out=denoised_arr)

return denoised_arr