def _chi_fit_tv(dat, sigma=None, df=None, lam=10, max_iter=50, tol=1e-5):
noise, tissue = estimate_noise(dat, chi=True)
mu = tissue['mean']
sigma = sigma or noise['sd']
df = df or noise['dof']
print(f'sigma = {sigma}, dof = {df}, mu = {mu}')
isigma2 = 1 / (sigma * sigma)
lam = lam / mu
msk = torch.isfinite(dat)
if msk.any():
dat = dat.masked_fill(msk.bitwise_not_(), 0)
fit = chi_bias_correction(dat, sigma, df)[0]
msk = dat > 0
n = msk.sum()
w = torch.ones_like(dat)
h = w.new_full([1] * dat.dim(), isigma2)[None]
ll_prev = float('inf')
for n_iter in range(max_iter):
w, tv = spatial.membrane_weights(fit[None],
factor=lam,
return_sum=True)
l, delta = nll_chi(dat, fit, msk, isigma2, df)
delta += spatial.regulariser(fit[None], membrane=lam, weights=w)[0]
delta = spatial.solve_field(h, delta[None], membrane=lam, weights=w)[0]
fit.sub_(delta).clamp_min_(1e-8 * sigma)
ll = l + tv
gain, ll_prev = ll_prev - ll, ll
print(f'{n_iter:3d} | {l/n:12.6g} + {tv/n:12.6g} = {ll/n:12.6g} '
f'| gain = {gain/n:6.3}')
if abs(gain) < tol * n:
break
return fit, sigma, df
def do_inpaint(dat,
voxel_size=1,
max_iter_rls=50,
max_iter_cg=32,
tol_rls=1e-5,
tol_cg=1e-5,
verbose=1):
"""Perform inpainting
Parameters
----------
dat : (channels, *spatial) tensor
Tensor with missing data marked as NaN
Returns
-------
dat : (channels, *spatial) tensor
"""
dat = torch.as_tensor(dat)
backend = utils.backend(dat)
voxel_size = utils.make_vector(voxel_size, dat.dim() - 1, **backend)
voxel_size = voxel_size / voxel_size.mean()
# initialize
mask = torch.isnan(dat)
for channel in dat:
channel[torch.isnan(channel)] = channel[~torch.isnan(channel)].median()
weights = dat.new_ones([1, *dat.shape[1:]])
ll_w = weights.sum(dtype=torch.double)
ll_x = spatial.membrane(dat, dim=3, voxel_size=voxel_size, weights=weights)
ll_x = (ll_x * dat).sum(dtype=torch.double)
if verbose > 0:
print(f'{0:3d} | {ll_x} + {ll_w} = {ll_x + ll_w}')
# Reweighted least squares loop
zeros = torch.zeros_like(dat)
for n_iter_rls in range(1, max_iter_rls + 1):
ll_x0 = ll_x
ll_w0 = ll_w
grad = spatial.membrane(dat,
dim=3,
voxel_size=voxel_size,
weights=weights)
grad = grad.masked_select(mask)
def precond(x):
p = spatial.membrane_diag(dim=3,
voxel_size=voxel_size,
weights=weights)
p = p.expand_as(dat).masked_select(mask)
p = p
return x / p
def forward(x):
m = zeros
m.masked_scatter_(mask, x)
m = spatial.membrane(m,
dim=3,
voxel_size=voxel_size,
weights=weights)
m = m.masked_select(mask)
return m
delta = optim.cg(forward,
grad,
precond=precond,
max_iter=max_iter_cg,
tolerance=tol_cg,
verbose=verbose > 1,
stop='norm')
dat.masked_scatter_(mask, dat.masked_select(mask) - delta)
weights, ll_w = spatial.membrane_weights(dat,
dim=3,
voxel_size=voxel_size,
return_sum=True)
ll_x = spatial.membrane(dat,
dim=3,
voxel_size=voxel_size,
weights=weights)
ll_x = (ll_x * dat).sum(dtype=torch.double)
if verbose > 0:
print(f'{n_iter_rls:3d} | {ll_x} + {ll_w} = {ll_x + ll_w}')
if abs(((ll_x0 - ll_w0) - (ll_x + ll_w)) / (ll_x0 - ll_w0)) < tol_rls:
if verbose > 0:
print('Converged')
break
return dat
def greeq(data, transmit=None, receive=None, opt=None, **kwopt):
"""Fit a non-linear relaxometry model to multi-echo Gradient-Echo data.
Parameters
----------
data : sequence[GradientEchoMulti]
Observed GRE data.
transmit : sequence[PrecomputedFieldMap], optional
Map(s) of the transmit field (b1+). If a single map is provided,
it is used to correct all contrasts. If multiple maps are
provided, there should be one for each contrast.
receive : sequence[PrecomputedFieldMap], optional
Map(s) of the receive field (b1-). If a single map is provided,
it is used to correct all contrasts. If multiple maps are
provided, there should be one for each contrast.
If no receive map is provided, the output `pd` map will have
a remaining b1- bias field.
opt : GREEQOptions or dict, optional
Algorithm options.
{'preproc': {'register': True}, # Co-register contrasts
'optim': {'nb_levels': 1, # Number of pyramid levels
'max_iter_rls': 10, # Max reweighting iterations
'max_iter_gn': 5, # Max Gauss-Newton iterations
'max_iter_cg': 32, # Max Conjugate-Gradient iterations
'tolerance': 1e-05, # Tolerance for early stopping (RLS)
'tolerance': 1e-05, ""
'tolerance_cg': 1e-03}, ""
'backend': {'dtype': torch.float32, # Data type
'device': 'cpu'}, # Device
'penalty': {'norm': 'jtv', # Type of penalty: {'tkh', 'tv', 'jtv', None}
'factor': {'r1': 10, # Penalty factor per (log) map
'pd': 10,
'r2s': 2,
'mt': 2}},
'verbose': 1}
Returns
-------
pd : ParameterMap
Proton density
r1 : ParameterMap
Longitudinal relaxation rate
r2s : ParameterMap
Apparent transverse relaxation rate
mt : ParameterMap, optional
Magnetisation transfer saturation
Only returned is MT-weighted data is provided.
"""
opt = GREEQOptions().update(opt, **kwopt)
dtype = opt.backend.dtype
device = opt.backend.device
backend = dict(dtype=dtype, device=device)
# --- estimate noise / register / initialize maps ---
data, transmit, receive, maps = preproc(data, transmit, receive, opt)
has_mt = hasattr(maps, 'mt')
# --- prepare penalty factor ---
lam = opt.penalty.factor
if isinstance(lam, dict):
lam = [
lam.get('pd', 0),
lam.get('r1', 0),
lam.get('r2s', 0),
lam.get('mt', 0)
]
if not has_mt:
lam = lam[:3]
lam = core.utils.make_vector(lam, 3 + has_mt,
**backend) # PD, R1, R2*, [MT]
# --- initialize weights (RLS) ---
if str(opt.penalty.norm).lower().startswith('no') or all(lam == 0):
opt.penalty.norm = ''
opt.penalty.norm = opt.penalty.norm.lower()
mean_shape = maps[0].shape
rls = None
sumrls = 0
if opt.penalty.norm in ('tv', 'jtv'):
rls_shape = mean_shape
if opt.penalty.norm == 'tv':
rls_shape = (len(maps), ) + rls_shape
rls = torch.ones(rls_shape, **backend)
sumrls = 0.5 * core.py.prod(rls_shape)
if opt.penalty.norm:
print(f'With {opt.penalty.norm.upper()} penalty:')
print(f' - PD: {lam[0]:.3g}')
print(f' - R1: {lam[1]:.3g}')
print(f' - R2*: {lam[2]:.3g}')
if has_mt:
print(f' - MT: {lam[3]:.3g}')
else:
print('Without penalty')
if opt.penalty.norm not in ('tv', 'jtv'):
# no reweighting -> do more gauss-newton updates instead
opt.optim.max_iter_gn *= opt.optim.max_iter_rls
opt.optim.max_iter_rls = 1
print('Optimization:')
print(f' - Tolerance: {opt.optim.tolerance}')
if opt.penalty.norm.endswith('tv'):
print(f' - IRLS iterations: {opt.optim.max_iter_rls}')
print(f' - GN iterations: {opt.optim.max_iter_gn}')
if opt.optim.solver == 'fmg':
print(f' - FMG cycles: 2')
print(f' - CG iterations: 2')
else:
print(f' - CG iterations: {opt.optim.max_iter_cg}'
f' (tolerance: {opt.optim.tolerance_cg})')
if opt.optim.nb_levels > 1:
print(f' - Levels: {opt.optim.nb_levels}')
printer = CritPrinter(max_levels=opt.optim.nb_levels,
max_rls=opt.optim.max_iter_rls,
max_gn=opt.optim.max_iter_gn,
penalty=opt.penalty.norm,
verbose=opt.verbose)
printer.print_head()
shape0 = shape = maps.shape[1:]
aff0 = aff = maps.affine
vx0 = vx = spatial.voxel_size(aff0)
vol0 = vx0.prod()
vol = vx.prod() / vol0
ll_scl = sum(core.py.prod(dat.shape) for dat in data)
for level in range(opt.optim.nb_levels, 0, -1):
printer.level = level
if opt.optim.nb_levels > 1:
aff, shape = _get_level(level, aff0, shape0)
vx = spatial.voxel_size(aff)
vol = vx.prod() / vol0
maps, rls = resize(maps, rls, aff, shape)
if opt.penalty.norm in ('tv', 'jtv'):
sumrls = 0.5 * vol * rls.reciprocal().sum(dtype=torch.double)
# --- compute derivatives ---
nb_prm = len(maps)
nb_hes = nb_prm * (nb_prm + 1) // 2
grad = torch.empty((nb_prm, ) + shape, **backend)
hess = torch.empty((nb_hes, ) + shape, **backend)
ll_rls = []
ll_gn = []
ll_max = float('inf')
max_iter_rls = max(opt.optim.max_iter_rls // level, 1)
for n_iter_rls in range(max_iter_rls):
# --- Reweighted least-squares loop ---
printer.rls = n_iter_rls
# --- Gauss Newton loop ---
for n_iter_gn in range(opt.optim.max_iter_gn):
# start = time.time()
printer.gn = n_iter_gn
crit = 0
grad.zero_()
hess.zero_()
# --- loop over contrasts ---
for contrast, b1m, b1p in zip(data, receive, transmit):
crit1, g1, h1 = derivatives_parameters(
contrast, maps, b1m, b1p, opt)
# increment
if hasattr(maps, 'mt') and not contrast.mt:
# we optimize for mt but this particular contrast
# has no information about mt so g1/h1 are smaller
# than grad/hess.
grad[:-1] += g1
hind = list(range(nb_prm - 1))
cnt = nb_prm
for i in range(nb_prm):
for j in range(i + 1, nb_prm):
if i != nb_prm - 1 and j != nb_prm - 1:
hind.append(cnt)
cnt += 1
hess[hind] += h1
crit += crit1
else:
grad += g1
hess += h1
crit += crit1
del g1, h1, crit1
# duration = time.time() - start
# print('grad', duration)
# start = time.time()
reg = 0.
if opt.penalty.norm:
g = spatial.regulariser(maps.volume,
weights=rls,
dim=3,
voxel_size=vx,
membrane=1,
factor=lam * vol)
grad += g
reg = 0.5 * dot(maps.volume, g)
del g
# duration = time.time() - start
# print('reg', duration)
# --- gauss-newton ---
# start = time.time()
grad = check_nans_(grad, warn='gradient')
hess = check_nans_(hess, warn='hessian')
if opt.penalty.norm:
# hess = hessian_sym_loaddiag(hess, 1e-5, 1e-8) # 1e-5 1e-8
if opt.optim.solver == 'fmg':
deltas = spatial.solve_field_fmg(hess,
grad,
rls,
factor=lam * vol,
membrane=1,
voxel_size=vx,
nb_iter=2)
else:
deltas = spatial.solve_field(
hess,
grad,
rls,
factor=lam * vol,
membrane=1,
voxel_size=vx,
verbose=max(0, opt.verbose - 1),
optim='cg',
max_iter=opt.optim.max_iter_cg,
tolerance=opt.optim.tolerance_cg,
stop='diff')
else:
# hess = hessian_sym_loaddiag(hess, 1e-3, 1e-4)
deltas = spatial.solve_field_closedform(hess, grad)
deltas = check_nans_(deltas, warn='deltas')
# duration = time.time() - start
# print('solve', duration)
for map, delta in zip(maps, deltas):
map.volume -= delta
map.volume.clamp_(-64, 64) # avoid exp overflow
del delta
del deltas
# --- Compute gain ---
ll = crit + reg + sumrls
ll_max = max(ll_max, ll)
ll_prev = ll_gn[-1] if ll_gn else ll_max
gain = ll_prev - ll
ll_gn.append(ll)
printer.print_crit(crit, reg, sumrls, gain / ll_scl)
if gain < opt.optim.tolerance * ll_scl:
# print('GN converged: ', ll_prev.item(), '->', ll.item())
break
# --- Update RLS weights ---
if opt.penalty.norm in ('tv', 'jtv'):
rls, sumrls = spatial.membrane_weights(
maps.volume,
lam,
vx,
return_sum=True,
dim=3,
joint=opt.penalty.norm == 'jtv',
eps=core.constants.eps(rls.dtype))
sumrls *= 0.5 * vol
# --- Compute gain ---
# (we are late by one full RLS iteration when computing the
# gain but we save some computations)
ll = ll_gn[-1]
ll_prev = ll_rls[-1] if ll_rls else ll_max
ll_rls.append(ll)
gain = ll_prev - ll
if abs(gain) < opt.optim.tolerance * ll_scl:
# print(f'RLS converged ({gain:7.2g})')
break
del grad
if opt.uncertainty:
uncertainty = compute_uncertainty(hess, rls, lam * vol, vx, opt)
maps.pd.uncertainty = uncertainty[0]
maps.r1.uncertainty = uncertainty[1]
maps.r2s.uncertainty = uncertainty[2]
if hasattr(maps, 'mt'):
maps.mt.uncertainty = uncertainty[3]
# --- Prepare output ---
return postproc(maps)
def correct_smooth(x,
sigma=None,
lam=10,
gamma=10,
downsample=None,
max_iter=16,
max_rls=8,
tol=1e-6,
verbose=False,
device=None):
"""Correct the intensity non-uniformity in a SPIM image.
The signal is modelled as: f = exp(s + b) + eps, with a penalty on
the (Squared) gradients of s and on the (squared) curvature of b.
Parameters
----------
x : tensor
SPIM image with the z dimension last and the z=0 plane first
sigma : float, optional
Noise standard deviation. Default: educated guess.
lam : float, default=10
Regularisation on the signal.
gamma : float, default=10
Regularisation on the bias field.
max_iter : int, default=16
Maximum number of Newton iterations.
max_rls : int, default=8
Maximum number of reweighting iterations.
If 1, this is effectively an l2 regularisation.
tol : float, default=1e-6
Tolerance for early stopping.
verbose : int or bool, default=False
Verbosity level
device : torch.device, default=x.device
Use this device during fitting.
Returns
-------
y : tensor
Fitted image
bias : float
Fitted bias
x : float
Corrected image
"""
x = torch.as_tensor(x)
if not x.dtype.is_floating_point:
x = x.to(dtype=torch.get_default_dtype())
dim = x.dim()
# downsampling
if downsample:
x0 = x
downsample = py.make_list(downsample, dim)
x = spatial.pool(dim, x, downsample)
shape = x.shape
x = x.to(device)
# noise educated guess: assume SNR=5 at z=1/2
center = tuple(slice(s // 3, 2 * s // 3) for s in shape)
sigma = sigma or x[center].median() / 5
lam = lam**2 * sigma**2
gamma = gamma**2 * sigma**2
regy = lambda y, w: spatial.regulariser(
y[None], membrane=lam, dim=dim, weights=w)[0]
regb = lambda b: spatial.regulariser(b[None], bending=gamma, dim=dim)[0]
solvey = lambda h, g, w: spatial.solve_field_sym(
h[None], g[None], membrane=lam, dim=dim, weights=w)[0]
solveb = lambda h, g: spatial.solve_field_sym(
h[None], g[None], bending=gamma, dim=dim)[0]
# init
l1 = max_rls > 1
if l1:
w = torch.ones_like(x)[None]
llw = w.sum()
max_rls = 10
else:
w = None
llw = 0
max_rls = 1
logb = torch.zeros_like(x)
logy = x.clamp_min(1e-3).log_()
y = logy.exp()
b = logb.exp()
fit = y * b
res = fit - x
llx = res.square().sum()
lly = (regy(logy, w).mul_(logy)).sum()
llb = (regb(logb).mul_(logb)).sum()
ll0 = llx + lly + llb + llw
ll1 = ll0
for it_ls in range(max_rls):
for it in range(max_iter):
# update bias
g = h = fit
h = (h * res).abs_()
h.addcmul_(g, g)
g *= res
g += regb(logb)
logb -= solveb(h, g)
logb0 = logb.mean()
logb -= logb0
logy += logb0
# update fit / ll
llb = (regb(logb).mul_(logb)).sum()
b = torch.exp(logb, out=b)
y = torch.exp(logy, out=y)
fit = y * b
res = fit - x
# update y
g = h = fit
h = (h * res).abs_()
h.addcmul_(g, g)
g *= res
g += regy(logy, w)
logy -= solvey(h, g, w)
# update fit / ll
y = torch.exp(logy, out=y)
fit = y * b
res = fit - x
lly = (regy(logy, w).mul_(logy)).sum()
# compute objective
llx = res.square().sum()
ll = llx + lly + llb + llw
gain = (ll1 - ll) / ll0
ll1 = ll
if verbose:
end = '\n' if verbose > 1 else '\r'
pre = f'{it_ls:3d} | ' if l1 else ''
print(pre + f'{it:3d} | {ll:12.6g} | gain = {gain:12.6g}',
end=end)
if it > 0 and abs(gain) < tol:
break
if l1:
w, llw = spatial.membrane_weights(logy[None],
lam,
dim=dim,
return_sum=True)
ll0 = ll
if verbose:
print('')
if downsample:
b = spatial.resize(logb.to(x0.device)[None, None],
downsample,
shape=x0.shape,
anchor='f')[0, 0].exp_()
y = spatial.resize(logy.to(x0.device)[None, None],
downsample,
shape=x0.shape,
anchor='f')[0, 0].exp_()
x = x0
else:
y = torch.exp(logy, out=y)
x = x / b
return y, b, x
def run_epic_noreverse(echoes,
voxshift=None,
lam=1,
sigma=1,
max_iter=(10, 32),
tol=1e-5,
verbose=False):
"""Run EPIC on pre-loaded tensors (no reverse gradients or synthesis)
Parameters
----------
echoes : (N, *spatial) tensor
Echoes acquired with bipolar readout, Readout direction should be last.
lam : [list of] float
Regularization factor (per echo)
sigma : float
Noise standard deviation
max_iter : [pair of] int
Maximum number of RLS and CG iterations
tol : float
Tolerance for early stopping
verbose : int,
Verbosity level
Returns
-------
echoes : (N, *spatial) tensor
Undistorted + denoised echoes
"""
ne = len(echoes) # number of echoes
nv = echoes.shape[1:].numel() # number of voxels
nd = echoes.dim() - 1 # number of dimensions
# initialize denoised echoes
fit = echoes.clone()
fwd_fit = torch.empty_like(fit)
# prepare voxel shift maps
do_voxshift = voxshift is not None
if do_voxshift:
ivoxshift = add_identity_1d(-voxshift)
voxshift = add_identity_1d(voxshift)
else:
ivoxshift = None
# prepare parameters
max_iter, sub_iter = py.make_list(max_iter, 2)
tol, sub_tol = py.make_list(tol, 2)
lam = [l / ne for l in py.make_list(lam, ne)]
isigma2 = 1 / (sigma * sigma)
# compute hessian once and for all
if do_voxshift:
one = torch.empty_like(voxshift)[None]
h = torch.empty_like(echoes)
h[0::2] = push1d(pull1d(one, voxshift), voxshift)
h[1::2] = push1d(pull1d(one, ivoxshift), ivoxshift)
del one
else:
h = fit.new_ones([1] * (nd + 1))
h *= isigma2
loss = float('inf')
for n_iter in range(max_iter):
# update weights
w, jtv = membrane_weights(fit, factor=lam, return_sum=True)
# gradient of likelihood
pull_forward(fit, voxshift, ivoxshift, out=fwd_fit)
fwd_fit.sub_(echoes)
ll = ssq(fwd_fit)
push_forward(fwd_fit, voxshift, ivoxshift, out=fwd_fit)
g = fwd_fit.mul_(isigma2)
ll *= 0.5 * isigma2
# gradient of prior
g += regulariser(fit, membrane=1, factor=lam, weights=w)
# solve
fit -= solve_field(h,
g,
w,
membrane=1,
factor=lam,
max_iter=sub_iter,
tolerance=sub_tol)
# track objective
ll, jtv = ll.item() / (ne * nv), jtv.item() / (ne * nv)
loss, loss_prev = ll + jtv, loss
if n_iter:
gain = (loss_prev - loss) / max((loss_max - loss), 1e-8)
else:
gain = float('inf')
loss_max = loss
if verbose:
end = '\n' if verbose > 1 else '\r'
print(
f'{n_iter+1:02d} | {ll:12.6g} + {jtv:12.6g} = {loss:12.6g} '
f'| gain = {gain:12.6g}',
end=end)
if gain < tol:
break
if verbose == 1:
print('')
return fit
def run_epic(echoes,
reverse_echoes=None,
voxshift=None,
extrapolate=True,
lam=1,
sigma=1,
max_iter=(10, 32),
tol=1e-5,
verbose=False):
"""Run EPIC on pre-loaded tensors.
Parameters
----------
echoes : (N, *spatial) tensor
Echoes acquired with bipolar readout, Readout direction should be last.
reverse_echoes : (N, *spatial)
Echoes acquired with reverse bipolar readout. Else: synthesized.
voxshift : (*spatial) tensor
Voxel shift map used to deform towards even (0, 2, ...) echoes.
Its inverse is used to deform towards odd (1, 3, ...) echoes.
extrapolate : bool
Extrapolate first/last echo when reverse_echoes is None.
Otherwise, only use interpolated echoes.
lam : [list of] float
Regularization factor (per echo)
sigma : float
Noise standard deviation
max_iter : [pair of] int
Maximum number of RLS and CG iterations
tol : float
Tolerance for early stopping
verbose : int,
Verbosity level
Returns
-------
echoes : (N, *spatial) tensor
Undistorted + denoised echoes
"""
if reverse_echoes is False:
return run_epic_noreverse(echoes, voxshift, lam, sigma, max_iter, tol,
verbose)
ne = len(echoes) # number of echoes
nv = echoes.shape[1:].numel() # number of voxels
nd = echoes.dim() - 1 # number of dimensions
# synthesize echoes
synth = not torch.is_tensor(reverse_echoes)
if synth:
neg = synthesize_neg(echoes[0::2])
pos = synthesize_pos(echoes[1::2])
reverse_echoes = torch.stack([x for y in zip(pos, neg) for x in y])
del pos, neg
else:
extrapolate = True
# initialize denoised echoes
fit = (echoes + reverse_echoes).div_(2)
if not extrapolate:
fit[0] = echoes[0]
fit[-1] = echoes[-1]
fwd_fit = torch.zeros_like(fit)
bwd_fit = torch.zeros_like(fit)
# prepare voxel shift maps
if voxshift is not None:
ivoxshift = add_identity_1d(-voxshift)
voxshift = add_identity_1d(voxshift)
else:
ivoxshift = None
# prepare parameters
max_iter, sub_iter = py.make_list(max_iter, 2)
tol, sub_tol = py.make_list(tol, 2)
lam = [l / ne for l in py.make_list(lam, ne)]
isigma2 = 1 / (sigma * sigma)
# compute hessian once and for all
if voxshift is not None:
one = torch.ones_like(voxshift)[None]
if extrapolate:
h = push1d(pull1d(one, voxshift), voxshift)
h += push1d(pull1d(one, ivoxshift), ivoxshift)
weight_ = lambda x: x.mul_(0.5)
halfweight_ = lambda x: x.mul_(math.sqrt(0.5))
else:
h = torch.zeros_like(fit)
h[:-1] += push1d(pull1d(one, voxshift), voxshift)
h[1:] += push1d(pull1d(one, ivoxshift), ivoxshift)
weight_ = lambda x: x[1:-1].mul_(0.5)
halfweight_ = lambda x: x[1:-1].mul_(math.sqrt(0.5))
del one
weight_(h)
else:
h = fit.new_ones([ne] + [1] * nd)
if extrapolate:
h *= 2
weight_ = lambda x: x.mul_(0.5)
halfweight_ = lambda x: x.mul_(math.sqrt(0.5))
else:
h[1:-1] *= 2
weight_ = lambda x: x[1:-1].mul_(0.5)
halfweight_ = lambda x: x[1:-1].mul_(math.sqrt(0.5))
weight_(h)
h *= isigma2
loss = float('inf')
for n_iter in range(max_iter):
# update weights
w, jtv = membrane_weights(fit, factor=lam, return_sum=True)
# gradient of likelihood (forward)
pull_forward(fit, voxshift, ivoxshift, out=fwd_fit)
fwd_fit.sub_(echoes)
halfweight_(fwd_fit)
ll = ssq(fwd_fit)
halfweight_(fwd_fit)
push_forward(fwd_fit, voxshift, ivoxshift, out=fwd_fit)
# gradient of likelihood (reversed)
pull_backward(fit, voxshift, ivoxshift, extrapolate, out=bwd_fit)
if extrapolate:
bwd_fit.sub_(reverse_echoes)
else:
bwd_fit[1:-1].sub_(reverse_echoes[1:-1])
halfweight_(bwd_fit)
ll += ssq(bwd_fit)
halfweight_(bwd_fit)
push_backward(bwd_fit, voxshift, ivoxshift, extrapolate, out=bwd_fit)
g = fwd_fit.add_(bwd_fit).mul_(isigma2)
ll *= 0.5 * isigma2
# gradient of prior
g += regulariser(fit, membrane=1, factor=lam, weights=w)
# solve
fit -= solve_field(h,
g,
w,
membrane=1,
factor=lam,
max_iter=sub_iter,
tolerance=sub_tol)
# track objective
ll, jtv = ll.item() / (ne * nv), jtv.item() / (ne * nv)
loss, loss_prev = ll + jtv, loss
if n_iter:
gain = (loss_prev - loss) / max((loss_max - loss), 1e-8)
else:
gain = float('inf')
loss_max = loss
if verbose:
end = '\n' if verbose > 1 else '\r'
print(
f'{n_iter+1:02d} | {ll:12.6g} + {jtv:12.6g} = {loss:12.6g} '
f'| gain = {gain:12.6g}',
end=end)
if gain < tol:
break
if verbose == 1:
print('')
return fit
def denoise(image=None,
lam=1,
sigma=1,
max_iter=64,
sub_iter=32,
optim='cg',
tol=1e-5,
sub_tol=1e-5,
plot=False,
jtv=True,
dim=None,
**prm):
"""Denoise an image using a (joint) total variation prior.
This implementation uses a reweighted least squares approach.
Parameters
----------
image : (..., K, *spatial) tensor, Image to denoise with K channels
lam : [list of] float, default=1, Regularisation factor
max_iter : int, default=64, Number of RLS iterations
sub_iter : int, default=32, Number of relaxation/cg iterations
optim : {'cg', 'relax', 'fmg+cg', 'fmg+relax'}, default='cg'
plot : bool, default=False
jtv : bool, default=True, Joint TV across channels ($\ell_{1,2}$)
dim : int, default=image.dim()-1
Returns
-------
denoised : (..., K, *spatial) tensor, Denoised image
"""
if image is None:
torch.random.manual_seed(1234)
image = phantoms.augment(phantoms.circle(), fwhm=0)[None]
image = torch.as_tensor(image)
dim = dim or (image.dim() - 1)
nvox = image.shape[-dim:].numel()
# regularization (prior)
lam = make_list(lam, image.shape[-dim - 1])
if jtv:
lam = [l / image.shape[-dim - 1] for l in lam]
prm['membrane'] = 1
prm['factor'] = lam
# noise variance (likelihood)
sigma = make_vector(sigma,
image.shape[-dim - 1],
dtype=image.dtype,
device=image.device)
isigma2 = 1 / (sigma * sigma)
# solver
optim, lr = make_list(optim, 2, default=1)
do_fmg, optim = ('fmg' in optim), ('relax' if 'relax' in optim else 'cg')
if do_fmg:
def solve(h, g, w, optim):
return spatial.solve_field_fmg(h,
g,
dim=dim,
weights=w,
**prm,
optim=optim)
else:
def solve(h, g, w, optim):
return spatial.solve_field(h,
g,
dim=dim,
weights=w,
**prm,
optim=optim,
max_iter=sub_iter,
tolerance=sub_tol)
# initialize
denoised = image.clone()
l_prev = None
l_max = None
for n_iter in range(1, max_iter + 1):
# update weight map / tv loss
weights, lw = spatial.membrane_weights(denoised,
dim=dim,
return_sum=True,
joint=jtv,
factor=lam)
# gradient / hessian of least square problem
ll, g, h = losses.mse(denoised, image, dim=dim, lam=isigma2)
ll.mul_(nvox), g.mul_(nvox), h.mul_(nvox)
g += spatial.regulariser(denoised, dim=dim, weights=weights, **prm)
# solve least square problem
optim0 = 'cg' if n_iter < 10 else optim # it seems it helps
g = solve(h, g, weights, optim0)
if lr != 1:
g.mul_(lr)
denoised -= g
ll = ll.item() / image.numel()
lw = lw.item() / image.numel()
l = ll + lw
# print objective
if l_prev is None:
l_prev = l_max = l
gain = float('inf')
print(f'{n_iter:03d} | {ll:12.6g} + {lw:12.6g} = {l:12.6g}',
end='\r')
else:
gain = (l_prev - l) / max(abs(l_max - l), 1e-8)
l_prev = l
print(
f'{n_iter:03d} | {ll:12.6g} + {lw:12.6g} = {l:12.6g} '
f'| gain = {gain:12.6g}',
end='\r')
if plot and (n_iter - 1) % (max_iter // 10 + 1) == 0:
import matplotlib.pyplot as plt
img = image[0, image.shape[1] // 2] if dim == 3 else image[0]
den = denoised[0,
denoised.shape[1] // 2] if dim == 3 else denoised[0]
wgt = weights[0, weights.shape[1] // 2] if dim == 3 else weights[0]
plt.subplot(1, 3, 1)
plt.imshow(img.cpu())
plt.subplot(1, 3, 2)
plt.imshow(den.cpu())
plt.subplot(1, 3, 3)
plt.imshow(wgt.reciprocal().cpu())
plt.colorbar()
plt.show()
if gain < tol:
break
print('')
return denoised
def meetup(data, dist=None, opt=None):
"""Fit the ESTATICS+MEETUP model to multi-echo Gradient-Echo data.
Parameters
----------
data : sequence[GradientEchoMulti]
Observed GRE data.
dist : sequence[Optional[ParameterizedDistortion]], optional
Pre-computed distortion fields
opt : Options, optional
Algorithm options.
Returns
-------
intecepts : sequence[GradientEcho]
Echo series extrapolated to TE=0
decay : estatics.ParameterMap
R2* decay map
distortions : sequence[ParameterizedDistortion]
B0-induced distortion fields
"""
# --- prepare everything -------------------------------------------
data, maps, dist, opt, rls = _prepare(data, dist, opt)
sumrls = 0.5 * rls.sum(dtype=torch.double) if rls is not None else 0
backend = dict(dtype=opt.backend.dtype, device=opt.backend.device)
# --- initialize tracking of the objective function ----------------
crit, vreg, reg = float('inf'), 0, 0
ll_rls = crit + vreg + reg
sumrls_prev = sumrls
rls_changed = False
ll_scl = sum(core.py.prod(dat.shape) for dat in data)
# --- Multi-Resolution loop ----------------------------------------
shape0 = shape = maps.shape[1:]
aff0 = aff = maps.affine
lam0 = lam = opt.regularization.factor
vx0 = vx = spatial.voxel_size(aff0)
scl0 = vx0.prod()
armijo_prev = 1
for level in range(opt.optim.nb_levels, 0, -1):
if opt.optim.nb_levels > 1:
aff, shape = _get_level(level, aff0, shape0)
vx = spatial.voxel_size(aff)
scl = vx.prod() / scl0
lam = [float(l * scl) for l in lam0]
maps, rls = _resize(maps, rls, aff, shape)
if opt.regularization.norm in ('tv', 'jtv'):
sumrls = 0.5 * scl * rls.reciprocal().sum(dtype=torch.double)
grad = torch.empty((len(data) + 1, *shape), **backend)
hess = torch.empty((len(data) * 2 + 1, *shape), **backend)
# --- RLS loop -------------------------------------------------
max_iter_rls = 1 if level > 1 else opt.optim.max_iter_rls
for n_iter_rls in range(1, max_iter_rls + 1):
# --- helpers ----------------------------------------------
regularizer_prm = lambda x: spatial.regulariser(
x, weights=rls, dim=3, voxel_size=vx, membrane=1, factor=lam)
solve_prm = lambda H, g: spatial.solve_field(
H,
g,
rls,
factor=lam,
membrane=1 if opt.regularization.norm else 0,
voxel_size=vx,
max_iter=opt.optim.max_iter_cg,
tolerance=opt.optim.tolerance_cg,
dim=3)
reweight = lambda x, **k: spatial.membrane_weights(
x,
lam,
vx,
dim=3,
**k,
joint=opt.regularization.norm == 'jtv',
eps=core.constants.eps(rls.dtype))
# ----------------------------------------------------------
# Initial update of parameter maps
# ----------------------------------------------------------
crit_pre_prm = None
max_iter_prm = opt.optim.max_iter_prm
for n_iter_prm in range(1, max_iter_prm + 1):
crit_prev = crit
reg_prev = reg
# --- loop over contrasts ------------------------------
crit = 0
grad.zero_()
hess.zero_()
for i, (contrast, intercept, distortion) \
in enumerate(zip(data, maps.intercepts, dist)):
# compute gradient
crit1, g1, h1 = derivatives_parameters(
contrast, distortion, intercept, maps.decay, opt)
# increment
gind, hind = [i, -1], [i, len(grad) - 1, len(grad) + i]
grad[gind] += g1
hess[hind] += h1
crit += crit1
del g1, h1
# --- regularization -----------------------------------
reg = 0.
if opt.regularization.norm:
g1 = regularizer_prm(maps.volume)
reg = 0.5 * dot(maps.volume, g1)
grad += g1
del g1
if n_iter_prm == 1:
crit_pre_prm = crit
# --- track RLS improvement ----------------------------
# Updating the RLS weights changes `sumrls` and `reg`. Now
# that we have the updated value of `reg` (before it is
# modified by the map update), we can check that updating the
# weights indeed improved the objective.
if opt.verbose and rls_changed:
rls_changed = False
if reg + sumrls <= reg_prev + sumrls_prev:
evol = '<='
else:
evol = '>'
ll_tmp = crit_prev + reg + vreg + sumrls
pstr = (f'{n_iter_rls-1:3d} | {"---":3s} | {"rls":4s} | '
f'{crit_prev:12.6g} + {reg:12.6g} + '
f'{sumrls:12.6g} + {vreg:12.6g} '
f'= {ll_tmp:12.6g} | {evol}')
if opt.optim.nb_levels > 1:
pstr = f'{level:3d} | ' + pstr
print(pstr)
# --- gauss-newton -------------------------------------
# Computing the GN step involves solving H\g
hess = check_nans_(hess, warn='hessian')
hess = hessian_loaddiag_(hess, 1e-6, 1e-8)
deltas = solve_prm(hess, grad)
deltas = check_nans_(deltas, warn='delta')
dd = regularizer_prm(deltas)
dv = dot(dd, maps.volume)
dd = dot(dd, deltas)
delta_reg = 0.5 * (dd - 2 * dv)
for map, delta in zip(maps, deltas):
map.volume -= delta
if map.min is not None or map.max is not None:
map.volume.clamp_(map.min, map.max)
del deltas
# --- track parameter map improvement --------------
gain = (crit_prev + reg_prev) - (crit + reg)
if n_iter_prm > 1 and gain < opt.optim.tolerance * ll_scl:
break
# ----------------------------------------------------------
# Distortion update with line search
# ----------------------------------------------------------
max_iter_dist = opt.optim.max_iter_dist
for n_iter_dist in range(1, max_iter_dist + 1):
crit = 0
vreg = 0
new_crit = 0
new_vreg = 0
deltas, dd, dv = [], 0, 0
# --- loop over contrasts ------------------------------
for i, (contrast, intercept, distortion) \
in enumerate(zip(data, maps.intercepts, dist)):
# --- helpers --------------------------------------
vxr = distortion.voxel_size[contrast.readout]
lam_dist = dict(opt.distortion.factor)
lam_dist['factor'] *= vxr**2
regularizer_dist = lambda x: spatial.regulariser(
x[None],
**lam_dist,
dim=3,
bound=DIST_BOUND,
voxel_size=distortion.voxel_size)[0]
solve_dist = lambda h, g: spatial.solve_field_fmg(
h,
g,
**lam_dist,
dim=3,
bound=DIST_BOUND,
voxel_size=distortion.voxel_size)
# --- likelihood -----------------------------------
crit1, g, h = derivatives_distortion(
contrast, distortion, intercept, maps.decay, opt)
crit += crit1
# --- regularization -------------------------------
vol = distortion.volume
g1 = regularizer_dist(vol)
vreg1 = 0.5 * vol.flatten().dot(g1.flatten())
vreg += vreg1
g += g1
del g1
# --- gauss-newton ---------------------------------
h = check_nans_(h, warn='hessian (distortion)')
h = hessian_loaddiag_(h[None], 1e-32, 1e-32, sym=True)[0]
delta = solve_dist(h, g)
delta = check_nans_(delta, warn='delta (distortion)')
deltas.append(delta)
del g, h
deltas.append(delta)
dd1 = regularizer_dist(delta)
dv += dot(dd1, vol)
dd1 = dot(dd1, delta)
dd += dd1
del delta, vol
# --- track parameters improvement ---------------------
if opt.verbose and n_iter_dist == 1:
gain = crit_pre_prm - (crit + delta_reg)
evol = '<=' if gain > 0 else '>'
ll_tmp = crit + reg + vreg + sumrls
pstr = (f'{n_iter_rls:3d} | {"---":3} | {"prm":4s} | '
f'{crit:12.6g} + {reg + delta_reg:12.6g} + '
f'{sumrls:12.6g} + {vreg:12.6g} '
f'= {ll_tmp:12.6g} | '
f'gain = {gain / ll_scl:7.2g} | {evol}')
if opt.optim.nb_levels > 1:
pstr = f'{level:3d} | ' + pstr
print(pstr)
# --- line search ----------------------------------
reg = reg + delta_reg
new_crit, new_reg = crit, reg
armijo, armijo_prev, ok = armijo_prev, 0, False
maps0 = maps
for n_ls in range(1, 12 + 1):
for delta1, dist1 in zip(deltas, dist):
dist1.volume.sub_(delta1, alpha=(armijo - armijo_prev))
armijo_prev = armijo
new_vreg = 0.5 * armijo * (armijo * dd - 2 * dv)
maps = maps0.deepcopy()
max_iter_prm = opt.optim.max_iter_prm
for n_iter_prm in range(1, max_iter_prm + 1):
crit_prev = new_crit
reg_prev = new_reg
# --- loop over contrasts ------------------
new_crit = 0
grad.zero_()
hess.zero_()
for i, (contrast, intercept, distortion) in enumerate(
zip(data, maps.intercepts, dist)):
# compute gradient
new_crit1, g1, h1 = derivatives_parameters(
contrast, distortion, intercept, maps.decay,
opt)
# increment
gind, hind = [i, -1
], [i,
len(grad) - 1,
len(grad) + i]
grad[gind] += g1
hess[hind] += h1
new_crit += new_crit1
del g1, h1
# --- regularization -----------------------
new_reg = 0.
if opt.regularization.norm:
g1 = regularizer_prm(maps.volume)
new_reg = 0.5 * dot(maps.volume, g1)
grad += g1
del g1
new_gain = (crit_prev + reg_prev) - (new_crit +
new_reg)
if new_gain < opt.optim.tolerance * ll_scl:
break
# --- gauss-newton -------------------------
# Computing the GN step involves solving H\g
hess = check_nans_(hess, warn='hessian')
hess = hessian_loaddiag_(hess, 1e-6, 1e-8)
delta = solve_prm(hess, grad)
delta = check_nans_(delta, warn='delta')
dd = regularizer_prm(delta)
dv = dot(dd, maps.volume)
dd = dot(dd, delta)
delta_reg = 0.5 * (dd - 2 * dv)
for map, delta1 in zip(maps, delta):
map.volume -= delta1
if map.min is not None or map.max is not None:
map.volume.clamp_(map.min, map.max)
del delta
if new_crit + new_reg + new_vreg <= crit + reg:
ok = True
break
else:
armijo = armijo / 2
if not ok:
for delta1, dist1 in zip(deltas, dist):
dist1.volume.add_(delta1, alpha=armijo_prev)
armijo_prev = 1
maps = maps0
new_crit = crit
new_vreg = 0
del delta
else:
armijo_prev *= 1.5
new_vreg = vreg + new_vreg
# --- track distortion improvement ---------------------
if not ok:
evol = '== (x)'
elif new_crit + new_vreg <= crit + vreg:
evol = f'<= ({n_ls:2d})'
else:
evol = f'> ({n_ls:2d})'
gain = (crit + vreg + reg) - (new_crit + new_vreg + new_reg)
crit, reg, reg = new_crit, new_vreg, new_reg
if opt.verbose:
ll_tmp = crit + reg + vreg + sumrls
pstr = (
f'{n_iter_rls:3d} | {n_iter_dist:3d} | {"dist":4s} | '
f'{crit:12.6g} + {reg:12.6g} + {sumrls:12.6g} ')
pstr += f'+ {vreg:12.6g} '
pstr += f'= {ll_tmp:12.6g} | gain = {gain / ll_scl:7.2g} | '
pstr += f'{evol}'
if opt.optim.nb_levels > 1:
pstr = f'{level:3d} | ' + pstr
print(pstr)
if opt.plot:
_show_maps(maps, dist, data)
if not ok:
break
if n_iter_dist > 1 and gain < opt.optim.tolerance * ll_scl:
break
# --------------------------------------------------------------
# Update RLS weights
# --------------------------------------------------------------
if level == 1 and opt.regularization.norm in ('tv', 'jtv'):
rls_changed = True
sumrls_prev = sumrls
rls, sumrls = reweight(maps.volume, return_sum=True)
sumrls = 0.5 * sumrls
# --- compute gain -----------------------------------------
# (we are late by one full RLS iteration when computing the
# gain but we save some computations)
ll_rls_prev = ll_rls
ll_rls = crit + reg + vreg
gain = ll_rls_prev - ll_rls
if gain < opt.optim.tolerance * ll_scl:
print(f'RLS converged ({gain / ll_scl:7.2g})')
break
# --- prepare output -----------------------------------------------
out = postproc(maps, data)
if opt.distortion.enable:
out = (*out, dist)
return out
def nonlin(data, dist=None, opt=None):
"""Fit the ESTATICS model to multi-echo Gradient-Echo data.
Parameters
----------
data : sequence[GradientEchoMulti]
Observed GRE data.
dist : sequence[Optional[ParameterizedDistortion]], optional
Pre-computed distortion fields
opt : Options, optional
Algorithm options.
Returns
-------
intecepts : sequence[GradientEcho]
Echo series extrapolated to TE=0
decay : estatics.ParameterMap
R2* decay map
distortions : sequence[ParameterizedDistortion], if opt.distortion.enable
B0-induced distortion fields
"""
if opt.distortion.enable:
return meetup(data, dist, opt)
# --- prepare everything -------------------------------------------
data, maps, dist, opt, rls = _prepare(data, dist, opt)
sumrls = 0.5 * rls.sum(dtype=torch.double) if rls is not None else 0
backend = dict(dtype=opt.backend.dtype, device=opt.backend.device)
# --- initialize tracking of the objective function ----------------
crit = float('inf')
vreg = 0
reg = 0
sumrls_prev = sumrls
rls_changed = False
ll_gn = ll_rls = float('inf')
ll_scl = sum(core.py.prod(dat.shape) for dat in data)
# --- Multi-Resolution loop ----------------------------------------
shape0 = shape = maps.shape[1:]
aff0 = aff = maps.affine
lam0 = lam = opt.regularization.factor
vx0 = vx = spatial.voxel_size(aff0)
scl0 = vx0.prod()
for level in range(opt.optim.nb_levels, 0, -1):
if opt.optim.nb_levels > 1:
aff, shape = _get_level(level, aff0, shape0)
vx = spatial.voxel_size(aff)
scl = vx.prod() / scl0
lam = [float(l * scl) for l in lam0]
maps, rls = _resize(maps, rls, aff, shape)
if opt.regularization.norm in ('tv', 'jtv'):
sumrls = 0.5 * scl * rls.reciprocal().sum(dtype=torch.double)
grad = torch.empty((len(data) + 1, *shape), **backend)
hess = torch.empty((len(data) * 2 + 1, *shape), **backend)
# --- RLS loop -------------------------------------------------
# > max_iter_rls == 1 if regularization is not (J)TV
max_iter_rls = opt.optim.max_iter_rls
if level != 1:
max_iter_rls = 1
for n_iter_rls in range(1, max_iter_rls + 1):
# --- helpers ----------------------------------------------
regularizer = lambda x: spatial.regulariser(
x, weights=rls, dim=3, voxel_size=vx, membrane=1, factor=lam)
solve = lambda h, g: spatial.solve_field(
h,
g,
rls,
factor=lam,
membrane=1 if opt.regularization.norm else 0,
voxel_size=vx,
max_iter=opt.optim.max_iter_cg,
tolerance=opt.optim.tolerance_cg,
dim=3)
reweight = lambda x, **k: spatial.membrane_weights(
x,
lam,
vx,
dim=3,
**k,
joint=opt.regularization.norm == 'jtv',
eps=core.constants.eps(rls.dtype))
# ----------------------------------------------------------
# Update parameter maps
# ----------------------------------------------------------
max_iter_prm = opt.optim.max_iter_prm
for n_iter_prm in range(1, max_iter_prm + 1):
crit_prev = crit
reg_prev = reg
# --- loop over contrasts ------------------------------
crit = 0
grad.zero_()
hess.zero_()
for i, (contrast, intercept, distortion) \
in enumerate(zip(data, maps.intercepts, dist)):
# compute gradient
crit1, g1, h1 = derivatives_parameters(
contrast, distortion, intercept, maps.decay, opt)
# increment
gind = [i, -1]
grad[gind] += g1
hind = [i, len(grad) - 1, len(grad) + i]
hess[hind] += h1
crit += crit1
del g1, h1
# --- regularization -----------------------------------
reg = 0.
if opt.regularization.norm:
g1 = regularizer(maps.volume)
reg = 0.5 * dot(maps.volume, g1)
grad += g1
del g1
# --- track RLS improvement ----------------------------
# Updating the RLS weights changes `sumrls` and `reg`. Now
# that we have the updated value of `reg` (before it is
# modified by the map update), we can check that updating the
# weights indeed improved the objective.
if opt.verbose and rls_changed:
rls_changed = False
if reg + sumrls <= reg_prev + sumrls_prev:
evol = '<='
else:
evol = '>'
ll_tmp = crit_prev + reg + vreg + sumrls
pstr = (
f'{n_iter_rls:3d} | {"---":3s} | {"rls":4s} | '
f'{crit_prev:12.6g} + {reg:12.6g} + {sumrls:12.6g} ')
if opt.distortion.enable:
pstr += f'+ {vreg:12.6g} '
pstr += f'= {ll_tmp:12.6g} | '
pstr += f'{evol}'
if opt.optim.nb_levels > 1:
pstr = f'{level:3d} | ' + pstr
print(pstr)
# --- gauss-newton -------------------------------------
# Computing the GN step involves solving H\g
hess = check_nans_(hess, warn='hessian')
hess = hessian_loaddiag_(hess, 1e-6, 1e-8)
deltas = solve(hess, grad)
deltas = check_nans_(deltas, warn='delta')
for map, delta in zip(maps, deltas):
map.volume -= delta
if map.min is not None or map.max is not None:
map.volume.clamp_(map.min, map.max)
del deltas
# --- compute GN gain ----------------------------------
ll_gn_prev = ll_gn
ll_gn = crit + reg + vreg + sumrls
gain = ll_gn_prev - ll_gn
if opt.verbose:
pstr = f'{n_iter_rls:3d} | {n_iter_prm:3d} | '
if opt.distortion.enable:
pstr += f'{"----":4s} | '
pstr += f'{"-"*72:72s} | '
else:
pstr += f'{"prm":4s} | '
pstr += f'{crit:12.6g} + {reg:12.6g} + {sumrls:12.6g} '
pstr += f'= {ll_gn:12.6g} | '
pstr += f'gain = {gain/ll_scl:7.2g}'
if opt.optim.nb_levels > 1:
pstr = f'{level:3d} | ' + pstr
print(pstr)
if opt.plot:
_show_maps(maps, dist, data)
if gain < opt.optim.tolerance * ll_scl:
break
# --------------------------------------------------------------
# Update RLS weights
# --------------------------------------------------------------
if level == 1 and opt.regularization.norm in ('tv', 'jtv'):
reg = 0.5 * dot(maps.volume, regularizer(maps.volume))
rls_changed = True
sumrls_prev = sumrls
rls, sumrls = reweight(maps.volume, return_sum=True)
sumrls = 0.5 * sumrls
# --- compute gain -----------------------------------------
# (we are late by one full RLS iteration when computing the
# gain but we save some computations)
ll_rls_prev = ll_rls
ll_rls = ll_gn
gain = ll_rls_prev - ll_rls
if gain < opt.optim.tolerance * ll_scl:
print(f'Converged ({gain/ll_scl:7.2g})')
break
# --- prepare output -----------------------------------------------
out = postproc(maps, data)
if opt.distortion.enable:
out = (*out, dist)
return out