def compose_polar_faces(faces, overlap, num_cores=None):
nside = faces[0].shape[0]
assert faces[0].shape[1] == nside
assert overlap <= nside // 4
# FIXME: If this is changed, change the constructor of PW_Transform!
nside_output = 2 * (3 * nside // 4 + overlap)
cut = nside // 4 - overlap
size = nside - cut
# north = np.ones((4, size, size))
north = sm.empty((4, size, size))
# north[0] = np.rot90(faces[1][:size, cut:], -1)
# north[1] = np.rot90(faces[2][:size, cut:], -2)
# north[2] = np.rot90(faces[0][:size, cut:], 0)
# north[3] = np.rot90(faces[3][:size, cut:], -3)
first_index = (1, 2, 0, 3)
second_index = (-1, -2, 0, -3)
with sharedmem_pool(num_cores, numexpr=False) as pool:
def work(i):
north[i] = np.rot90(faces[first_index[i]][:size, cut:],
second_index[i])
pool.map(work, range(4))
output = np.ones((2, nside_output, nside_output))
output[0] = np.vstack((np.hstack(
(north[0], north[1])), np.hstack((north[2], north[3]))))
del north
# south = np.ones((4, size, size))
south = sm.empty((4, size, size))
first_index = (9, 8, 10, 11)
second_index = (1, 0, 2, 3)
# south[0] = np.rot90(faces[9][cut:, :size], 1)
# south[1] = np.rot90(faces[8][cut:, :size], 0)
# south[2] = np.rot90(faces[10][cut:, :size], 2)
# south[3] = np.rot90(faces[11][cut:, :size], 3)
with sharedmem_pool(num_cores, numexpr=False) as pool:
def work(i):
south[i] = np.rot90(faces[first_index[i]][cut:, :size],
second_index[i])
pool.map(work, range(4))
output[1] = np.vstack((np.hstack(
(south[0], south[1])), np.hstack((south[2], south[3]))))
return output
def inv_thresh_scales(self, coeffs, noise_sigma=None, multipliers=None):
#INV THRESHOLD THE SCALES FOR A FACE
if noise_sigma is None:
print('ABORT: a sigma value must be specified')
return hp_im
assert noise_sigma >= 0.0
if multipliers is None:
multipliers = [3] * trafo.num_scales + [4]
width = self.__trafo.width
height = self.__trafo.height
indices = self.__trafo.indices
for kl, ksc in enumerate(indices):
if (kl == 0):
scales = [0]
else:
scales.append(ksc[0] + 1)
thresh_lvls = [noise_sigma * multipliers[sc] for sc in scales]
thresh_coeff = sm.empty_like(coeffs, dtype=DTYPES[1])
with sharedmem_pool(self.__num_cores) as pool:
def work(i):
coeff = coeffs[i]
thresh = thresh_lvls[i]
ne.evaluate('coeff*(real(abs(coeff))>=thresh)',
out=thresh_coeff[i])
pool.map(work, range(len(indices)))
thresh_coeff = np.array(thresh_coeff)
recon = self.__trafo.inverse_transform(thresh_coeff,
real=True,
do_norm=True)
return recon
def get_faces_par(im, nested=False, num_cores=None):
npix = np.shape(im)[0]
assert npix % 12 == 0
nside = npix2nside(npix)
if not nested:
new_im = reorder(im, r2n=True)
else:
new_im = im
# index = np.array([xyf2pix(nside, x, range(nside), 0, True)
# for x in range(nside-1, -1, -1)])
index = make_index_par(nside, num_cores)
CubeFace = sm.empty((12, nside, nside))
with sharedmem_pool(num_cores, numexpr=False) as pool:
# for face in range(12):
# CubeFace[face] = np.resize(new_im[index + nside**2 * face],
# (nside, nside))
def work(i):
CubeFace[i] = np.resize(new_im[index + nside**2 * i],
(nside, nside))
pool.map(work, range(12))
return np.array(CubeFace)
def make_index_par(nside, num_cores=None):
# original sequential code:
# index = np.array([xyf2pix(nside, x, range(nside), 0, True)
# for x in range(nside-1, -1, -1)])
index = sm.empty((nside, nside), dtype='int64')
y = np.arange(nside)
with sharedmem_pool(num_cores, numexpr=False) as pool:
# we want to divide range(nside) into N chunks of (almost) equal size
real_core_num = pool.np
N = 10 * real_core_num # this can be changed if desired!
chunk_size = 1 + nside // N # NOTE: N * chunk_size >= nside
def work(i):
lower = (N - (i + 1)) * chunk_size
upper = (N - i) * chunk_size
if lower < nside:
index[lower:upper] = [
xyf2pix(nside, x, y, 0, True)
for x in range(lower, min(nside, upper))
]
pool.map(work, range(N - 1, -1, -1))
return np.array(index[::-1])
def set_faces_par(CubeFace, nested=False, num_cores=None):
npix = np.size(CubeFace)
assert npix % 12 == 0
nside = npix2nside(npix)
index = make_index_par(nside, num_cores)
# index = np.array([xyf2pix(nside, x, range(nside), 0, True)
# for x in range(nside-1, -1, -1)])
imag = sm.empty((npix))
with sharedmem_pool(num_cores, numexpr=False) as pool:
def work(i):
for face in range(12):
imag[index + nside**2 * face] = np.resize(
CubeFace[face], (nside, nside))
#imag[index + nside**2 * face] = np.resize(CubeFace[face],
# (nside,nside))
pool.map(work, range(12))
imag = np.array(imag)
if not nested:
new_im = reorder(imag, n2r=True)
else:
new_im = imag
return new_im
def _blend_masks_par(nside, size, num_cores=None):
def _cos(x):
# return -0.5 * np.cos(np.pi*x) + 0.5
pi = np.pi
return ne.evaluate('-0.5 * cos(pi * x) + 0.5')
x, y = np.meshgrid(
np.linspace(1, nside, nside),
np.hstack([
np.zeros(3 * nside // 4 - size),
np.linspace(0, 1, size),
np.ones(nside // 4)
]))
blend_masks = sm.empty((8, nside, nside))
blend_masks[0] = _cos(y)
# blend_masks[1] = np.rot90(blend_masks[0], -1)
# blend_masks[2] = np.rot90(blend_masks[0], -2)
# blend_masks[3] = np.rot90(blend_masks[0], -3)
with sharedmem_pool(num_cores, numexpr=False) as pool:
def work(i):
blend_masks[i] = np.rot90(blend_masks[0], -i)
pool.map(work, range(1, 4))
# blend_masks[4] = blend_masks[0] * blend_masks[3]
# blend_masks[5] = blend_masks[1] * blend_masks[2]
# blend_masks[6] = (1 - blend_masks[0]) * (1 - blend_masks[1])
# blend_masks[7] = (1 - blend_masks[2]) * (1 - blend_masks[3])
first_index = (0, 1, 0, 2)
second_index = (3, 2, 1, 3)
with sharedmem_pool(num_cores) as pool:
def work(i):
if i in {4, 5}:
blend_masks[i] = (blend_masks[first_index[i - 4]] *
blend_masks[second_index[i - 4]])
elif i in {6, 7}:
blend_masks[i] = ((1 - blend_masks[first_index[i - 4]]) *
(1 - blend_masks[second_index[i - 4]]))
pool.map(work, range(4, 8))
return np.array(blend_masks)
def compose_equatorial_faces(faces, overlap, num_cores=None):
nside = faces[0].shape[0]
assert faces[0].shape[1] == nside
assert overlap <= 3 * nside // 4
extra = nside // 4 + overlap
nside_output = nside + 2 * extra
masks = _face_masks(extra, num_cores)
output = sm.empty((4, nside_output, nside_output))
with sharedmem_pool(num_cores, numexpr=False) as pool:
def work(i):
rows = []
temp = []
temp.append(faces[4 + (i + 1) % 4][-extra:, -extra:])
temp.append(faces[i][-extra:, :])
temp.append(
(masks[0] * np.rot90(faces[i][-extra:, -extra:], -1) +
masks[1] * np.rot90(faces[(i + 3) % 4][:extra, :extra], 1)) /
(masks[0] + masks[1]))
rows.append(np.hstack(temp))
temp = []
temp.append(faces[8 + i][:, -extra:])
temp.append(faces[4 + i])
temp.append(faces[(i + 3) % 4][:, :extra])
rows.append(np.hstack(temp))
temp = []
temp.append(
(masks[2] * np.rot90(faces[8 +
(i + 3) % 4][:extra, :extra], -1) +
masks[3] * np.rot90(faces[8 + i][-extra:, -extra:], 1)) /
(masks[2] + masks[3]))
temp.append(faces[8 + (i + 3) % 4][:extra, :])
temp.append(faces[4 + (i + 3) % 4][:extra, :extra])
rows.append(np.hstack(temp))
output[i] = np.vstack(rows)
pool.map(work, range(4))
return np.array(output)
def max_coeff(self, hp_im, do_norm=True):
maxima = [0 for j in range(6)]
for j, face in enumerate(self._face_generator(hp_im)):
face_fourier = self.__trafo.do_fourier(face)
numexpr_flag = NUM_FFTW_THREADS != 1
with sharedmem_pool(self.__num_cores,
numexpr=numexpr_flag) as pool:
def work(i):
coeff = self.__trafo.transform_fourier(face_fourier,
i,
do_norm=do_norm)
return np.max(np.abs(coeff))
maxima[j] = np.max(
pool.map(work, range(self.__trafo.redundancy)))
return np.max(maxima)
def _face_masks(size, num_cores=None):
def _cos_step(x, size):
pi = np.pi
return ne.evaluate('(0 <= x) * (-0.5 * cos(pi*x/size) + 0.5)'
'* (x < size) + (x >= size)')
# return ((0 <= x) * (-0.5 * np.cos(np.pi*x/size) + 0.5) * (x < size)
# + (x >= size))
# masks = np.zeros((4, size, size))
masks = sm.empty((4, size, size))
# one of the values is shifted to avoid for the sum of two masks to vanish
x, y = np.meshgrid(np.linspace(0, size - 1, size),
np.linspace(1, size, size))
# masks[0] = np.rot90(_cos_step(size-x, size) * _cos_step(2*y-x, size), 1)
# masks[1] = np.rot90(_cos_step(size-y+1, size) * _cos_step(2*x-y+2, size), 1)
# masks[2] = np.rot90(masks[0], 2)
# masks[3] = np.rot90(masks[1], 2)
with sharedmem_pool(num_cores, numexpr=True) as pool:
def work(i):
if i == 0:
masks[0] = np.rot90(
_cos_step(size - x, size) * _cos_step(2 * y - x, size), 1)
masks[2] = np.rot90(masks[0], 2)
else:
masks[1] = np.rot90(
_cos_step(size - y + 1, size) *
_cos_step(2 * x - y + 2, size), 1)
masks[3] = np.rot90(masks[1], 2)
pool.map(work, range(2))
return np.array(masks)
def reassemble_faces_par(nside,
polar_faces,
eq_faces,
margin,
transition,
num_cores=None):
out = sm.empty((12, nside, nside))
blend_masks = _blend_masks_par(nside, transition, num_cores)
offset = 3 * nside // 4 + transition + margin
size = offset - margin
with sharedmem_pool(num_cores, numexpr=False) as pool:
def work(i):
if i <= 3:
### north polar faces
out[i][:size, -size:] = np.rot90(polar_faces[0],
i)[offset:-margin,
margin:offset]
out[i] *= blend_masks[6]
else:
### south polar faces
j = i - 8
out[i][-size:, :size] = np.rot90(polar_faces[1],
-j)[margin:offset,
offset:-margin]
out[i] *= blend_masks[7]
pool.map(work, [0, 1, 2, 3, 8, 9, 10, 11])
extra = nside // 4 + transition
offset = extra + margin
with sharedmem_pool(num_cores, numexpr=False) as pool:
def work(i):
if i < 4:
# north
out[i][-extra:] += (blend_masks[0][-extra:] *
eq_faces[i][margin:offset, offset:-offset])
# east
out[i][:, :extra] += (
blend_masks[1][:, :extra] *
eq_faces[(i + 1) % 4][offset:-offset, -offset:-margin])
elif i < 8:
# north-west
out[i][-extra:, -extra:] += (
blend_masks[4][-extra:, -extra:] *
eq_faces[(i - 1) % 4][margin:offset, margin:offset])
# middle
out[i] += (eq_faces[i - 4][offset:-offset, offset:-offset])
# south-east
out[i][:extra, :extra] += (
blend_masks[5][:extra, :extra] *
eq_faces[(i - 3) % 4][-offset:-margin, -offset:-margin])
else:
# west
out[i][:, -extra:] += (
blend_masks[3][:, -extra:] *
eq_faces[i - 8][offset:-offset, margin:offset])
# south
out[i][:extra] += (
blend_masks[2][:extra] *
eq_faces[(i - 7) % 4][-offset:-margin, offset:-offset])
pool.map(work, range(12))
with sharedmem_pool(num_cores, numexpr=False) as pool:
def work(i):
if i < 4:
out[i] /= (blend_masks[6] + blend_masks[0] + blend_masks[1])
elif i < 8:
out[i] /= (1 + blend_masks[4] + blend_masks[5])
else:
out[i] /= (blend_masks[7] + blend_masks[2] + blend_masks[3])
pool.map(work, range(12))
return np.array(out)