Source code for mrinufft.trajectories.gradients
"""Functions to improve/modify gradients."""
import numpy as np
import numpy.linalg as nl
from scipy.interpolate import CubicSpline
[docs]
def patch_center_anomaly(
shot_or_params,
update_shot=None,
update_parameters=None,
in_out=False,
learning_rate=1e-1,
):
"""Re-position samples to avoid center anomalies.
Some trajectories behave slightly differently from expected when
approaching definition bounds, most often the k-space center as
for spirals in some cases.
This function enforces non-strictly increasing monoticity of
sample distances from the center, effectively reducing slew
rates and smoothing gradient transitions locally.
Shots can be updated with provided functions to keep fitting
their strict definition, or it can be smoothed using cubic
spline approximations.
Parameters
----------
shot_or_params : np.array, list
Either a single shot of shape (Ns, Nd), or a list of arbitrary
arguments used by ``update_shot`` to initialize a single shot.
update_shot : function, optional
Function used to initialize a single shot based on parameters
provided by ``update_parameters``. If None, cubic splines are
used as an approximation instead, by default None
update_parameters : function, optional
Function used to update shot parameters when provided in
``shot_or_params`` from an updated shot and parameters.
If None, cubic spline parameterization is used instead,
by default None
in_out : bool, optional
Whether the shot is going in-and-out or start from the center,
by default False
learning_rate : float, optional
Learning rate used in the iterative optimization process,
by default 1e-1
Returns
-------
array_like
N-D trajectory based on ``shot_or_params`` if a shot or
update_shot otherwise.
list
Updated parameters either in the ``shot_or_params`` format
if params, or cubic spline parameterization as an array of
floats between 0 and 1 otherwise.
"""
if update_shot is None or update_parameters is None:
single_shot = shot_or_params
else:
parameters = shot_or_params
single_shot = update_shot(*parameters)
Ns = single_shot.shape[0]
old_index = 0
old_x_axis = np.zeros(Ns)
x_axis = np.linspace(0, 1, Ns)
if update_shot is None or update_parameters is None:
def _default_update_parameters(shot, *parameters):
return parameters
update_parameters = _default_update_parameters
parameters = [x_axis]
update_shot = CubicSpline(x_axis, single_shot)
while nl.norm(x_axis - old_x_axis) / Ns > 1e-10:
# Determine interval to fix
single_shot = update_shot(*parameters)
gradient_norm = nl.norm(
np.diff(single_shot[(Ns // 2) * in_out :], axis=0), axis=-1
)
arc_length = np.cumsum(gradient_norm)
index = gradient_norm[1:] - arc_length[1:] / (
1 + np.arange(1, len(gradient_norm))
)
index = np.where(index < 0, np.inf, index)
index = max(old_index, 2 + np.argmin(index))
start_index = (Ns // 2 + (Ns % 2) - index) * in_out
final_index = (Ns // 2) * in_out + index
# Balance over interval
cbs = CubicSpline(x_axis, single_shot.T, axis=1)
gradient_norm = nl.norm(
np.diff(single_shot[start_index:final_index], axis=0), axis=-1
)
old_x_axis = np.copy(x_axis)
new_x_axis = np.cumsum(np.diff(x_axis[start_index:final_index]) / gradient_norm)
new_x_axis *= (x_axis[final_index - 1] - x_axis[start_index]) / new_x_axis[
-1
] # Normalize
new_x_axis += x_axis[start_index]
x_axis[start_index + 1 : final_index - 1] += (
new_x_axis[:-1] - x_axis[start_index + 1 : final_index - 1]
) * learning_rate
# Update loop variables
old_index = index
single_shot = cbs(x_axis).T
parameters = update_parameters(single_shot, *parameters)
single_shot = single_shot = update_shot(*parameters)
return single_shot, parameters