def __init__(self, psfs, rho, **kwds):
super(ADMMLasso, self).__init__(**kwds)
self.shape = psfs.shape
# Initialize C library.
self.c_admm_lasso = admm_lasso.initialize3D(rho, self.shape[0],
self.shape[1],
self.shape[2])
#
# Do the ADMM math on the Python side.
#
# Calculate A matrix.
nz = self.shape[2]
A = admmMath.Cells(nz, 1)
for i in range(nz):
tmp = recenterPSF.recenterPSF(psfs[:, :, i])
A[i, 0] = numpy.fft.fft2(tmp)
# Calculate A transpose.
At = admmMath.transpose(A)
# Calculate AtA + rhoI inverse.
G = admmMath.multiplyMatMat(At, A)
for i in range(nz):
G[i, i] += admmMath.identityMatrix(G.getMatrixShape(), scale=rho)
[L, D, U] = admmMath.lduG(G)
L_inv = admmMath.invL(L)
D_inv = admmMath.invD(D)
U_inv = admmMath.invU(U)
G_inv = admmMath.multiplyMatMat(U_inv,
admmMath.multiplyMatMat(D_inv, L_inv))
# Initialize A and G_inv matrices in the C library.
fft_size = int(self.shape[1] / 2 + 1)
for i in range(nz):
# Remove redundant frequencies that FFTW doesn't use.
c_A = A[i, 0][:, :fft_size]
c_A = numpy.ascontiguousarray(c_A, dtype=numpy.complex128)
admm_lasso.initializeA(self.c_admm_lasso, c_A, i)
for i in range(nz):
for j in range(nz):
# Remove redundant frequencies that FFTW doesn't use.
#
# We index (j,i) here because this is what gives us results that
# match admm_3d (the pure Python version of 3D ADMM.
#
c_G_inv = G_inv[j, i][:, :fft_size]
c_G_inv = numpy.ascontiguousarray(c_G_inv,
dtype=numpy.complex128)
admm_lasso.initializeGInv(self.c_admm_lasso, c_G_inv,
i * nz + j)
def invU(U):
"""
Calculate inverse of U Cell.
"""
nr, nc = U.getCellsShape()
mshape = U.getMatrixShape()
assert (nr == nc), "U Cell must be square!"
nmat = nr
u_tmp = admmMath.copyCell(U)
u_inv = admmMath.Cells(nmat, nmat)
for i in range(nmat):
for j in range(nmat):
if (i == j):
u_inv[i,j] = numpy.identity(mshape[0])
else:
u_inv[i,j] = numpy.zeros_like(U[0,0])
for j in range(nmat-1,0,-1):
for i in range(j-1,-1,-1):
tmp = u_tmp[i,j]
for k in range(nmat):
u_tmp[i,k] = u_tmp[i,k] - numpy.matmul(tmp, u_tmp[j,k])
u_inv[i,k] = u_inv[i,k] - numpy.matmul(tmp, u_inv[j,k])
return u_inv
def invL(L):
"""
Calculate inverse of L Cell.
"""
nr, nc = L.getCellsShape()
mshape = L.getMatrixShape()
assert (nr == nc), "L Cell must be square!"
nmat = nr
l_tmp = admmMath.copyCell(L)
l_inv = admmMath.Cells(nmat, nmat)
for i in range(nmat):
for j in range(nmat):
if (i == j):
l_inv[i,j] = numpy.identity(mshape[0])
else:
l_inv[i,j] = numpy.zeros_like(L[0,0])
for j in range(nmat-1):
for i in range(j+1,nmat):
tmp = l_tmp[i,j]
for k in range(nmat):
l_tmp[i,k] = l_tmp[i,k] - numpy.matmul(tmp, l_tmp[j,k])
l_inv[i,k] = l_inv[i,k] - numpy.matmul(tmp, l_inv[j,k])
return l_inv
def lduG(G):
"""
G a Cell object containing AtA + rhoI matrices. The A
matrices are the PSF matrices, I is the identity matrix and
rho is the ADMM timestep.
"""
nr, nc = G.getCellsShape()
mshape = G.getMatrixShape()
assert (nr == nc), "G Cell must be square!"
nmat = nr
# Create empty M matrix.
M = admmMath.Cells(nmat, nmat)
for i in range(nmat):
for j in range(nmat):
M[i,j] = numpy.zeros_like(G[0,0])
# Schur decomposition.
D = admmMath.Cells(nmat, nmat)
L = admmMath.Cells(nmat, nmat)
U = admmMath.Cells(nmat, nmat)
for r in range(nmat-1,-1,-1):
for c in range(nmat-1,-1,-1):
k = max(r,c)
M[r,c] = G[r,c]
for s in range(nmat-1,k,-1):
M[r,c] = M[r,c] - numpy.matmul(M[r,s], numpy.matmul(numpy.linalg.inv(M[s,s]), M[s,c]))
if (r == c):
D[r,c] = M[r,c]
L[r,c] = numpy.identity(mshape[0])
U[r,c] = numpy.identity(mshape[0])
elif (r > c):
D[r,c] = numpy.zeros(mshape)
L[r,c] = numpy.matmul(M[r,c], numpy.linalg.inv(M[k,k]))
U[r,c] = numpy.zeros(mshape)
elif (r < c):
D[r,c] = numpy.zeros(mshape)
L[r,c] = numpy.zeros(mshape)
U[r,c] = numpy.matmul(numpy.linalg.inv(M[k,k]), M[r,c])
return [L, D, U]
def transpose(A):
"""
Returns the transpose of Cell.
"""
nr, nc = A.getCellsShape()
B = admmMath.Cells(nc, nr)
for r in range(nr):
for c in range(nc):
B[c,r] = numpy.transpose(A[r,c])
return B
def invD(D):
"""
Calculate inverse of D Cell.
"""
nr, nc = D.getCellsShape()
assert (nr == nc), "D Cell must be square!"
nmat = nr
d_inv = admmMath.Cells(nmat,nmat)
for i in range(nmat):
for j in range(nmat):
if (i == j):
d_inv[i,j] = numpy.linalg.inv(D[i,j])
else:
d_inv[i,j] = numpy.zeros_like(D[0,0])
return d_inv
def multiplyMatMat(A, B):
"""
Multiply two Cell objects following matrix multiplication rules.
"""
nr_a, nc_a = A.getCellsShape()
nr_b, nc_b = B.getCellsShape()
assert(nr_b == nc_a), "Cell sizes don't match!"
C = admmMath.Cells(nr_b, nc_a)
for r in range(nr_b):
for c in range(nc_a):
C[r,c] = numpy.zeros_like(A[0,0])
for k in range(nr_a):
C[r,c] += numpy.matmul(A[k,c], B[r,k])
return C
def __init__(self, psfs, rho, **kwds):
"""
psfs is an array of psfs for different image planes (nx, ny, nz).
They must be the same size as the image that will get analyzed.
"""
super(ADMM, self).__init__(**kwds)
self.A = None
self.At = None
self.Atb = None
self.G_inv = None
self.rho = rho
self.shape = psfs.shape
self.u = None
self.z = None
# Calculate A matrix.
nz = self.shape[2]
self.A = admmMath.Cells(nz, 1)
for i in range(nz):
tmp = recenterPSF.recenterPSF(psfs[:, :, i])
self.A[i, 0] = numpy.fft.fft2(tmp)
# Calculate A transpose.
self.At = admmMath.transpose(self.A)
# Calculate AtA + rhoI inverse.
G = admmMath.multiplyMatMat(self.At, self.A)
for i in range(nz):
G[i, i] += admmMath.identityMatrix(G.getMatrixShape(), scale=rho)
[L, D, U] = admmMath.lduG(G)
L_inv = admmMath.invL(L)
D_inv = admmMath.invD(D)
U_inv = admmMath.invU(U)
self.G_inv = admmMath.multiplyMatMat(
U_inv, admmMath.multiplyMatMat(D_inv, L_inv))