def test_pylops(xp, G, b):
G = xp.asarray(G)
b = xp.asarray(b)
svx = xp.linalg.solve(G, b)
svx2 = xp.tile(svx, 2)
G2 = pylops.BlockDiag([G, G])
b2 = xp.tile(b, 2)
xo, *_ = lsqr(G2, b2, show=show, atol=tol, btol=tol, iter_lim=maxit)
assert xp.allclose(xo, svx2, atol=tol, rtol=tol)
import numpy as np
import matplotlib.pyplot as plt
import pylops
plt.close('all')
###############################################################################
# Let's start by creating N MatrixMult operators and the BlockDiag operator
N = 100
Nops = 32
Ops = [
pylops.MatrixMult(np.random.normal(0., 1., (N, N))) for _ in range(Nops)
]
Op = pylops.BlockDiag(Ops, nproc=1)
###############################################################################
# We can now perform a scalability test on the forward operation
workers = [2, 3, 4]
compute_times, speedup = \
pylops.utils.multiproc.scalability_test(Op, np.ones(Op.shape[1]),
workers=workers, forward=True)
plt.figure(figsize=(12, 3))
plt.plot(workers, speedup, 'ko-')
plt.xlabel('# Workers')
plt.ylabel('Speed Up')
plt.title('Forward scalability test')
plt.tight_layout()
###############################################################################
# to three different subset of the model and data
#
# .. math::
# \mathbf{D_{BDiag}} =
# \begin{bmatrix}
# \mathbf{D_v} & \mathbf{0} & \mathbf{0} \\
# \mathbf{0} & 0.5*\mathbf{D_v} & \mathbf{0} \\
# \mathbf{0} & \mathbf{0} & -\mathbf{D_h}
# \end{bmatrix}, \qquad
# \mathbf{y} =
# \begin{bmatrix}
# \mathbf{D_v} \mathbf{x_1} \\
# 0.5*\mathbf{D_v} \mathbf{x_2} \\
# -\mathbf{D_h} \mathbf{x_3}
# \end{bmatrix}
BD = pylops.BlockDiag([D2vop, 0.5 * D2vop, -1 * D2hop])
Y = np.reshape(BD * X.ravel(), (3 * Nv, Nh))
fig, axs = plt.subplots(1, 2, figsize=(10, 3))
fig.suptitle("Block-diagonal", fontsize=14, fontweight="bold", y=0.95)
im = axs[0].imshow(X, interpolation="nearest")
axs[0].axis("tight")
axs[0].set_title(r"$x$")
plt.colorbar(im, ax=axs[0])
im = axs[1].imshow(Y, interpolation="nearest")
axs[1].axis("tight")
axs[1].set_title(r"$y$")
plt.colorbar(im, ax=axs[1])
plt.tight_layout()
plt.subplots_adjust(top=0.8)
def BlockDiagonalOperator(*linops: LinearOperator,
n_jobs: int = 1) -> PyLopLinearOperator:
r"""
Construct a block diagonal operator from N linear operators.
Parameters
----------
linops: LinearOperator
Linear operators forming the diagonal blocks.
Alternatively, numpy.ndarray or scipy.sparse.spmatrix can be passed in place of one or more operators.
n_jobs: int
Number of processes used to evaluate the N operators in parallel using multiprocessing.
If ``n_jobs=1`` (default), work in serial mode.
Returns
-------
PyLopLinearOperator
Block diagonal linear operator.
Examples
--------
.. doctest::
>>> from pycsou.linop.base import BlockDiagonalOperator
>>> from pycsou.linop.diff import SecondDerivative
>>> Nv, Nh = 11, 21
>>> D2hop = SecondDerivative(size=Nv * Nh, shape=(Nv,Nh), axis=1)
>>> D2vop = SecondDerivative(size=Nv * Nh, shape=(Nv,Nh), axis=0)
>>> Dblockdiag = BlockDiagonalOperator(D2vop, 0.5 * D2vop, -1 * D2hop)
>>> x = np.zeros((Nv, Nh)); x[int(Nv//2), int(Nh//2)] = 1; z = np.tile(x, (3,1)).flatten()
>>> np.allclose(Dblockdiag(z), np.concatenate((D2vop(x.flatten()), 0.5 * D2vop(x.flatten()), - D2hop(x.flatten()))))
True
Notes
-----
A block-diagonal operator composed of N linear operators is created such
as its application in forward mode leads to
.. math::
\begin{bmatrix}
\mathbf{L_1} & \mathbf{0} & \cdots & \mathbf{0} \\
\mathbf{0} & \mathbf{L_2} & \cdots & \mathbf{0} \\
\vdots & \vdots & \ddots & \vdots \\
\mathbf{0} & \mathbf{0} & \cdots & \mathbf{L_N}
\end{bmatrix}
\begin{bmatrix}
\mathbf{x}_{1} \\
\mathbf{x}_{2} \\
\vdots \\
\mathbf{x}_{N}
\end{bmatrix} =
\begin{bmatrix}
\mathbf{L_1} \mathbf{x}_{1} \\
\mathbf{L_2} \mathbf{x}_{2} \\
\vdots \\
\mathbf{L_N} \mathbf{x}_{N}
\end{bmatrix}
while its application in adjoint mode leads to
.. math::
\begin{bmatrix}
\mathbf{L_1}^\ast & \mathbf{0} & \cdots & \mathbf{0} \\
\mathbf{0} & \mathbf{L_2}^\ast & \cdots & \mathbf{0} \\
\vdots & \vdots & \ddots & \vdots \\
\mathbf{0} & \mathbf{0} & \cdots & \mathbf{L_N}^\ast
\end{bmatrix}
\begin{bmatrix}
\mathbf{y}_{1} \\
\mathbf{y}_{2} \\
\vdots \\
\mathbf{y}_{N}
\end{bmatrix} =
\begin{bmatrix}
\mathbf{L_1}^\ast \mathbf{y}_{1} \\
\mathbf{L_2}^\ast \mathbf{y}_{2} \\
\vdots \\
\mathbf{L_N}^\ast \mathbf{y}_{N}
\end{bmatrix}
The Lipschitz constant of the block-diagonal operator can be bounded by :math:`{\max_{i=1}^N \|\mathbf{L}_{i}\|_2}`.
Warnings
--------
The parameter ``n_jobs`` is currently unused and is there for compatibility with the future API of PyLops.
The code should be updated when the next version on PyLops is released.
See Also
--------
:py:class:`~pycsou.linop.base.BlockOperator`, :py:class:`~pycsou.linop.base.LinOpStack`
"""
pylinops = [linop.PyLop for linop in linops]
lipschitz_cst = np.array([linop.lipschitz_cst for linop in linops]).max()
block_diag = pylops.BlockDiag(ops=pylinops)
return PyLopLinearOperator(block_diag, lipschitz_cst=lipschitz_cst)
# \mathbf{0} & 0.5*\mathbf{D_v} & \mathbf{0} \\
# \mathbf{0} & \mathbf{0} & -\mathbf{D_h}
# \end{bmatrix}, \qquad
# \mathbf{y} =
# \begin{bmatrix}
# \mathbf{D_v} \mathbf{x_1} \\
# 0.5*\mathbf{D_v} \mathbf{x_2} \\
# -\mathbf{D_h} \mathbf{x_3}
# \end{bmatrix}
Nv, Nh = 11, 21
X = np.zeros((Nv * 3, Nh))
X[int(Nv / 2), int(Nh / 2)] = 1
X[int(Nv / 2) + Nv, int(Nh / 2)] = 1
X[int(Nv / 2) + 2 * Nv, int(Nh / 2)] = 1
Block = pylops.BlockDiag([D2vop, 0.5 * D2vop, -1 * D2hop])
Y = np.reshape(Block * np.ndarray.flatten(X), (11 * 3, 21))
fig, axs = plt.subplots(1, 2, figsize=(10, 3))
fig.suptitle('Block-diagonal', fontsize=14, fontweight='bold')
im = axs[0].imshow(X, interpolation='nearest')
axs[0].axis('tight')
axs[0].set_title(r'$x$')
plt.colorbar(im, ax=axs[0])
im = axs[1].imshow(Y, interpolation='nearest')
axs[1].axis('tight')
axs[1].set_title(r'$y$')
plt.colorbar(im, ax=axs[1])
plt.tight_layout()
def solve(self, f: np.ndarray):
"""
Description of main primal-dual iteration.
:return: None
"""
(primal_n, primal_m) = self.image_size
v = w = 0
g = f.ravel()
p = p_bar = np.zeros(primal_n*primal_m)
q = q_bar = np.zeros(2*primal_n*primal_m)
if self.reg_mode != 'tik':
grad = pylops.Gradient(dims=(primal_n, primal_m), dtype='float64', edge=True, kind="backward")
else:
grad = pylops.Identity(np.prod(self.image_size))
grad1 = pylops.BlockDiag([grad, grad]) # symmetric dxdy <-> dydx not necessary (expensive) but easy and functional
proj_0 = IndicatorL2((primal_n, primal_m), upper_bound=self.alpha[0])
proj_1 = IndicatorL2((2 * primal_n, primal_m), upper_bound=self.alpha[1])
if not self.silent_mode:
progress = progressbar.ProgressBar(max_value=self.max_iter)
k = 0
while (self.tol < self.sens or k == 0) and (k < self.max_iter):
p_old = p
q_old = q
# Dual Update
g = self.lam / (self.tau + self.lam) * (g + self.tau * (self.A*(p_bar ) - f)) #- self.alpha[0]*self.breg_p
if self.reg_mode != 'tik':
v = proj_0.prox(v + self.tau * (grad * p_bar - q_bar))
else:
v = self.alpha[0] / (self.tau + self.alpha[0]) * \
(v + self.tau * (grad*p_bar - self.data))
if self.reg_mode == 'tgv':
w = proj_1.prox(w + self.tau * grad1*q_bar)
# Primal Update
p = p - self.tau * (-self.alpha[0]*self.breg_p + self.A.H*g + grad.H * v)
if self.reg_mode == 'tgv':
q = q + self.tau * (v - grad1.H * w)
# Extragradient Update
p_bar = 2 * p - p_old
q_bar = 2 * q - q_old
if k % 50 == 0:
p_gap = p - p_old
self.sens = np.linalg.norm(p_gap)/np.linalg.norm(p_old)
print(self.sens)
if self.gamma:
raise NotImplementedError("The adjustment of the step size in the "
"Primal-Dual is not yet fully developed.")
thetha = 1 / np.sqrt(1 + 2*self.gamma * self.G.prox_param)
self.G.prox_param = thetha * self.G.prox_param
self.F_star.prox_param = self.F_star.prox_param / thetha
k += 1
if not self.silent_mode:
progress.update(k)
self.x = p
if k <= self.max_iter:
print(" Early stopping.")
return self.x
def trap_phase_2(image_vecs_medsub, model_vecs, temporal_basis,
trap_params: TrapParams):
xp = core.get_array_module(image_vecs_medsub)
was_gpu_array = xp is cp
timers = {}
flat_model_vecs = model_vecs.ravel()
if trap_params.scale_model_std:
model_coeff_scale = xp.std(flat_model_vecs)
flat_model_vecs /= model_coeff_scale
else:
model_coeff_scale = 1
if trap_params.force_gpu_fit:
temporal_basis = cp.asarray(temporal_basis)
flat_model_vecs = cp.asarray(flat_model_vecs)
operator_block_diag = [temporal_basis.T] * image_vecs_medsub.shape[0]
opstack = [
pylops.BlockDiag(operator_block_diag),
]
if trap_params.incorporate_offset:
if trap_params.background_split_mask is not None:
left_mask_vec = trap_params.background_split_mask
left_mask_megavec = np.repeat(left_mask_vec[:, np.newaxis],
model_vecs.shape[1]).ravel()
assert len(left_mask_megavec) == len(flat_model_vecs)
left_mask_megavec = left_mask_megavec[np.newaxis, :].astype(
flat_model_vecs.dtype)
left_mask_megavec = left_mask_megavec - left_mask_megavec.mean()
left_mask_megavec /= np.linalg.norm(left_mask_megavec)
# "ones" for left side pixels -> fit constant offset for left psf
opstack.append(xp.asarray(left_mask_megavec))
# "ones" for right side pixels -> fit constant offset for right psf
right_mask_megavec = -1 * left_mask_megavec
opstack.append(xp.asarray(right_mask_megavec))
else:
background_megavec = np.ones_like(flat_model_vecs[xp.newaxis, :])
background_megavec /= np.linalg.norm(background_megavec)
opstack.append(xp.asarray(background_megavec))
opstack.append(flat_model_vecs[xp.newaxis, :])
op = pylops.VStack(opstack).transpose()
log.debug(f"TRAP operator: {op}")
image_megavec = image_vecs_medsub.ravel()
log.debug(
f"Performing inversion on A.shape={op.shape} and b={image_megavec.shape}"
)
timers['invert'] = time.perf_counter()
if trap_params.use_cgls:
solver = pylops.optimization.solver.cgls
log.debug(f"{solver=}")
solver_kwargs = dict(damp=trap_params.damp, tol=trap_params.tol)
cgls_result = solver(op, image_megavec, xp.zeros(int(op.shape[1])),
**solver_kwargs)
xinv = cgls_result[0]
else:
soln = lsqr.lsqr(
op,
image_megavec,
x0=None,
atol=trap_params.tol,
damp=trap_params.damp,
# show=True,
# calc_var=True,
)
xinv = soln[0]
# import matplotlib.pyplot as plt
# plt.plot(soln[-1])
# plt.yscale('log')
timers['invert'] = time.perf_counter() - timers['invert']
log.debug(f"Finished RegularizedInversion in {timers['invert']} sec")
if core.get_array_module(xinv) is cp:
model_coeff = float(xinv.get()[-1])
else:
model_coeff = float(xinv[-1])
model_coeff = model_coeff / model_coeff_scale
# return model_coeff, timers
if trap_params.compute_residuals:
solnvec = xinv
solnvec[-1] = 0 # zero planet model contribution
log.debug(f"Constructing starlight estimate from fit vector")
timers['subtract'] = time.perf_counter()
estimate_vecs = op.dot(solnvec).reshape(image_vecs_medsub.shape)
if core.get_array_module(
image_vecs_medsub) is not core.get_array_module(estimate_vecs):
image_vecs_medsub = core.get_array_module(estimate_vecs).asarray(
image_vecs_medsub)
resid_vecs = image_vecs_medsub - estimate_vecs
if core.get_array_module(resid_vecs) is cp and not was_gpu_array:
resid_vecs = resid_vecs.get()
timers['subtract'] = time.perf_counter() - timers['subtract']
log.debug(f"Starlight subtracted in {timers['subtract']} sec")
return model_coeff, timers, resid_vecs
else:
return model_coeff, timers, None
