def _lens_chk_N_weave(args):
assert len(args) == 15, args
N, path_to_map, path_to_dx, path_to_dy, buff0, buff1, \
lside0, lside1, HD_res0, HD_res1, NR_iter, kspl, LD0, LD1, do_not_prefilter = args
HD_res = (HD_res0, HD_res1)
LD = (LD0, LD1)
s = (2**LD[0] + 2 * buff0, 2**LD[1] + 2 * buff1) # Chunk shape
assert s[0] == s[1], 'only square matrices here.' # TODO
dx = np.load(path_to_dx, mmap_mode='r')
dy = np.load(path_to_dy, mmap_mode='r')
map = np.load(path_to_map, mmap_mode='r')
rmin0 = lside0 / 2**HD_res[0]
rmin1 = lside1 / 2**HD_res[1]
filtmap = np.empty(s, dtype=np.float64)
dx_gu = np.empty(
s, dtype=np.float64
) # will dx displ. in grid units of each chunk (typ. (256 * 256) )
dy_gu = np.empty(
s, dtype=np.float64
) # will dy displ. in grid units of each chunk (typ. (256 * 256) )
lenmap = np.empty(s, dtype=np.float64)
sLDs, sHDs = periodicmap_spliter().get_slices_chk_N(
N, LD, HD_res, (buff0, buff1))
for sLD, sHD in zip(sLDs, sHDs):
# Displacements chunk in grid units, and map chunk to displace.
dx_gu[sLD] = dx[sHD] / rmin1
dy_gu[sLD] = dy[sHD] / rmin0
filtmap[sLD] = map[sHD]
if not do_not_prefilter:
# Prefiltering the map prior application of bicubic interpolation
filtmap = np.fft.rfft2(filtmap)
w0 = 6. / (2. * np.cos(2. * np.pi * np.fft.fftfreq(s[0])) + 4.)
filtmap *= np.outer(w0, w0[0:filtmap.shape[1]])
filtmap = np.fft.irfft2(filtmap, s)
bicubicspline = r"\
int i,j;\
for( j= 0; j < width; j++ )\
{\
for( i = 0; i < width; i++)\
{\
lenmap[j * width + i] = bicubiclensKernel(filtmap,i + dx_gu[j * width + i],j + dy_gu[j * width + i],width);\
}\
}"
width = int(s[0])
weave.inline(bicubicspline,
['lenmap', 'filtmap', 'dx_gu', 'dy_gu', 'width'],
headers=[bicubicspline_header])
return (lenmap, )
def get_dxdy_chk_N(self, N, buffers=None):
if buffers is None: buffers = self.buffers
shape = (2 ** self.LD_res[0] + 2 * buffers[0], 2 ** self.LD_res[1] + 2 * buffers[1])
dx, dy = np.zeros(shape), np.zeros(shape)
spliter_lib = map_spliter.periodicmap_spliter() # library to split periodic maps.
sLDs, sHDs = spliter_lib.get_slices_chk_N(N, self.LD_res, self.HD_res, buffers)
for sLD, sHD in zip(sLDs, sHDs):
dx[sLD] = self.get_dx()[sHD]
dy[sLD] = self.get_dy()[sHD]
return dx, dy
def _lens_chk_N(args):
assert len(args) == 15, args
N, path_to_map, path_to_dx, path_to_dy, buff0, buff1, \
lside0, lside1, HD_res0, HD_res1, NR_iter, kspl, LD0, LD1, do_not_prefilter = args
HD_res = (HD_res0, HD_res1)
LD = (LD0, LD1)
s = (2**LD[0] + 2 * buff0, 2**LD[1] + 2 * buff1) # Chunk shape
dx = np.load(path_to_dx, mmap_mode='r')
dy = np.load(path_to_dy, mmap_mode='r')
map = np.load(path_to_map, mmap_mode='r')
rmin0 = lside0 / 2**HD_res[0]
rmin1 = lside1 / 2**HD_res[1]
map_chk_N = np.empty(s, dtype=float)
dx_gu = np.empty(
s, dtype=float
) # will dx displ. in grid units of each chunk (typ. (256 * 256) )
dy_gu = np.empty(
s, dtype=float
) # will dy displ. in grid units of each chunk (typ. (256 * 256) )
sLDs, sHDs = periodicmap_spliter().get_slices_chk_N(
N, LD, HD_res, (buff0, buff1))
for sLD, sHD in zip(sLDs, sHDs):
# Displacements chunk in grid units, and map chunk to displace.
dx_gu[sLD] = dx[sHD] / rmin1
dy_gu[sLD] = dy[sHD] / rmin0
map_chk_N[sLD] = map[sHD]
if do_not_prefilter:
# Undoing the prefiltering prior to apply bicubic interpolation
map_chk_N = np.fft.rfft2(map_chk_N)
w0 = 6. / (2. * np.cos(2. * np.pi * np.fft.fftfreq(s[0])) + 4.)
map_chk_N /= np.outer(w0, w0[0:map_chk_N.shape[1]])
map_chk_N = np.fft.irfft2(map_chk_N, s)
idc0, idc1 = np.indices(
s) # Two typically (256 + 2 * buff * 256 + 2* buff) maps
lx = (idc1 + dx_gu).flatten() # No need to enforce periodicity here.
ly = (idc0 + dy_gu).flatten() # No need to enforce periodicity here.
return (interpolate.RectBivariateSpline(np.arange(s[0]),
np.arange(s[1]),
map_chk_N,
kx=kspl,
ky=kspl).ev(ly, lx).reshape(s), )
def Pool_generic(func, arg, root_Nthreads, do_not_prefilter=False):
assert len(arg) == 11, arg
path_to_map, path_to_dx, path_to_dy, buff0, buff1, lside0, lside1, HD_res0, HD_res1, NR_iter, kspl = arg
assert os.path.exists(path_to_dx) and os.path.exists(path_to_dy)
assert IsPowerOfTwo(root_Nthreads)
diff0, diff1 = Log2ofPowerof2((root_Nthreads, root_Nthreads))
HD_shape = (2**HD_res0, 2**HD_res1)
LD = (HD_res0 - diff0, HD_res1 - diff1)
pool = setup_Pool()
ret_list = pool.map(func, [[
i, path_to_map, path_to_dx, path_to_dy, buff0, buff1, lside0, lside1,
HD_res0, HD_res1, NR_iter, kspl, LD[0], LD[1], do_not_prefilter
] for i in range(root_Nthreads**2)])
pool.close()
pool.join()
# Recombines from the lensed_chks :
spliter_lib = periodicmap_spliter() # library to split periodic maps.
ret = [] # one map for lens, two for inverse
for i in range(len(ret_list[0])):
map = np.empty(HD_shape)
if verbose:
for j, N in enumerate_progress(
xrange(root_Nthreads**2),
label='Pool_generic:patching chks together'):
sLDs, sHDs = spliter_lib.get_slices_chk_N(N,
LD,
(HD_res0, HD_res1),
(buff0, buff1),
inverse=True)
map[sHDs[0]] = ret_list[N][i][sLDs[0]]
ret.append(map)
else:
for j, N in enumerate(xrange(root_Nthreads**2)):
sLDs, sHDs = spliter_lib.get_slices_chk_N(N,
LD,
(HD_res0, HD_res1),
(buff0, buff1),
inverse=True)
map[sHDs[0]] = ret_list[N][i][sLDs[0]]
ret.append(map)
return ret
def _get_inverse_chk_N_old(args):
"""
Returns inverse displacement in chunk N
Uses periodic boundary conditions, which is not applicable to chunks, thus there
will be boudary effects on the edges (2 or 4 pixels depending on the rule). Make sure the buffer is large enough.
"""
assert len(args) == 15, args
N, path_to_map, path_to_dx, path_to_dy, buff0, buff1, \
lside0, lside1, HD_res0, HD_res1, NR_iter, kspl, LD0, LD1, do_not_prefilter = args
HD_res = (HD_res0, HD_res1)
LD = (LD0, LD1)
s = (2**LD[0] + 2 * buff0, 2**LD[1] + 2 * buff1) # Chunk shape
rmin0 = lside0 / 2**HD_res[0]
rmin1 = lside1 / 2**HD_res[1]
# Get magn. matrix of the chunk:
extra_buff = np.array(
(5, 5)) # To avoid surprises with the periodic derivatives
dx = np.zeros(
s + 2 * extra_buff
) # will dx displ. in grid units of each chunk (typ. (256 * 256) )
dy = np.zeros(
s + 2 * extra_buff
) # will dy displ. in grid units of each chunk (typ. (256 * 256) )
sLDs, sHDs = periodicmap_spliter().get_slices_chk_N(
N, LD, HD_res, (buff0 + extra_buff[0], buff1 + extra_buff[1]))
for sLD, sHD in zip(sLDs, sHDs):
dx[sLD] = np.load(path_to_dx, mmap_mode='r')[sHD]
dy[sLD] = np.load(path_to_dy, mmap_mode='r')[sHD]
# Jacobian matrix of the chunk :
sl0 = slice(extra_buff[0], dx.shape[0] - extra_buff[0])
sl1 = slice(extra_buff[1], dx.shape[1] - extra_buff[1])
dfxdx_1 = PartialDerivativePeriodic(dx, axis=1, h=rmin1,
rule='4pts')[sl0, sl1] + 1.
dfydy_1 = PartialDerivativePeriodic(dy, axis=0, h=rmin0,
rule='4pts')[sl0, sl1] + 1.
dfydx = PartialDerivativePeriodic(dy, axis=1, h=rmin1, rule='4pts')[sl0,
sl1]
dfxdy = PartialDerivativePeriodic(dx, axis=0, h=rmin0, rule='4pts')[sl0,
sl1]
det = (dfxdx_1) * (dfydy_1) - dfydx * dfxdy
if not np.all(det > 0.):
print "ffs_displ::Negative value in det k : something's weird, you'd better check that"
# Inverse magn. elements. (with a minus sign) We may need to spline these later for further NR iterations :
_Minv_xx = -dfydy_1 / det
_Minv_yy = -dfxdx_1 / det
_Minv_xy = dfxdy / det
_Minv_yx = dfydx / det
del dfxdx_1, dfydx, dfydy_1, dfxdy
dx = dx[sl0, sl1] # Getting rid of extra buffer
dy = dy[sl0, sl1] # Getting rid of extra buffer
dxn = (_Minv_xx * dx + _Minv_xy * dy)
dyn = (_Minv_yx * dx + _Minv_yy * dy)
if NR_iter == 0: return dxn, dyn
# Setting up a bunch of splines to interpolate the increment to the displacement according to Newton-Raphson.
# Needed are splines of the forward displacement and of the (inverse, as implemented here) magnification matrix.
# Hopefully the map resolution is enough to spline the magnification matrix.
xcoord = np.arange(s[1]) * rmin1
ycoord = np.arange(s[0]) * rmin0
spl_dx = interpolate.RectBivariateSpline(ycoord,
xcoord,
dx,
kx=kspl,
ky=kspl)
spl_dy = interpolate.RectBivariateSpline(ycoord,
xcoord,
dy,
kx=kspl,
ky=kspl)
spl_xx = interpolate.RectBivariateSpline(ycoord,
xcoord,
_Minv_xx,
kx=kspl,
ky=kspl)
spl_yy = interpolate.RectBivariateSpline(ycoord,
xcoord,
_Minv_yy,
kx=kspl,
ky=kspl)
spl_xy = interpolate.RectBivariateSpline(ycoord,
xcoord,
_Minv_xy,
kx=kspl,
ky=kspl)
spl_yx = interpolate.RectBivariateSpline(ycoord,
xcoord,
_Minv_yx,
kx=kspl,
ky=kspl)
idc = np.indices(s)
y_x = idc[1] * rmin1
y_y = idc[0] * rmin0
del idc
for i in range(NR_iter):
dxn_1 = dxn
dyn_1 = dyn
lx = (y_x + dxn_1).flatten()
ly = (y_y + dyn_1).flatten()
res_x = dxn_1 + spl_dx.ev(ly, lx).reshape(s) # dx residuals
res_y = dyn_1 + spl_dy.ev(ly, lx).reshape(s) # dy residuals
dxn = dxn_1 + spl_xx.ev(ly, lx).reshape(s) * res_x + spl_xy.ev(
ly, lx).reshape(s) * res_y
dyn = dyn_1 + spl_yx.ev(ly, lx).reshape(s) * res_x + spl_yy.ev(
ly, lx).reshape(s) * res_y
return dxn, dyn
def _get_inverse_chk_N(args):
"""
Returns inverse displacement in chunk N
Uses periodic boundary conditions, which is not applicable to chunks, thus there
will be boudary effects on the edges (2 or 4 pixels depending on the rule). Make sure the buffer is large enough.
"""
assert len(args) == 15, args
N, path_to_map, path_to_dx, path_to_dy, buff0, buff1, \
lside0, lside1, HD_res0, HD_res1, NR_iter, kspl, LD0, LD1, do_not_prefilter = args
HD_res = (HD_res0, HD_res1)
LD = (LD0, LD1)
s = (2**LD[0] + 2 * buff0, 2**LD[1] + 2 * buff1) # Chunk shape
rmin0 = lside0 / 2**HD_res[0]
rmin1 = lside1 / 2**HD_res[1]
# Get magn. matrix of the chunk:
extra_buff = np.array(
(5, 5)) # To avoid surprises with the periodic derivatives
dx = np.zeros(
s + 2 * extra_buff
) # will dx displ. in grid units of each chunk (typ. (256 * 256) )
dy = np.zeros(
s + 2 * extra_buff
) # will dy displ. in grid units of each chunk (typ. (256 * 256) )
sLDs, sHDs = periodicmap_spliter().get_slices_chk_N(
N, LD, HD_res, (buff0 + extra_buff[0], buff1 + extra_buff[1]))
for sLD, sHD in zip(sLDs, sHDs):
dx[sLD] = np.load(
path_to_dx, mmap_mode='r'
)[sHD] / rmin1 # Need grid units displacement for the bicubic spline
dy[sLD] = np.load(path_to_dy, mmap_mode='r')[sHD] / rmin0
# Jacobian matrix of the chunk :
sl0 = slice(extra_buff[0], dx.shape[0] - extra_buff[0])
sl1 = slice(extra_buff[1], dx.shape[1] - extra_buff[1])
Minv_yy = -(PartialDerivativePeriodic(dx, axis=1)[sl0, sl1] + 1.)
Minv_xx = -(PartialDerivativePeriodic(dy, axis=0)[sl0, sl1] + 1.)
Minv_xy = PartialDerivativePeriodic(dy, axis=1)[sl0, sl1]
Minv_yx = PartialDerivativePeriodic(dx, axis=0)[sl0, sl1]
det = Minv_yy * Minv_xx - Minv_xy * Minv_yx
if not np.all(det > 0.):
print "ffs_displ::Negative value in det k : something's weird, you'd better check that"
# Inverse magn. elements. (with a minus sign) We may need to spline these later for further NR iterations :
Minv_xx /= det
Minv_yy /= det
Minv_xy /= det
Minv_yx /= det
del det
dx = dx[sl0, sl1] # Getting rid of extra buffer
dy = dy[sl0, sl1] # Getting rid of extra buffer
ex = (Minv_xx * dx + Minv_xy * dy)
ey = (Minv_yx * dx + Minv_yy * dy)
if NR_iter == 0: return ex * rmin1, ey * rmin0
# Setting up a bunch of splines to interpolate the increment to the displacement according to Newton-Raphson.
# Needed are splines of the forward displacement and of the (inverse, as implemented here) magnification matrix.
# Hopefully the map resolution is enough to spline the magnification matrix.
s0, s1 = dx.shape
r0 = s0
r1 = s1 / 2 + 1 # rfft shape
w0 = 6. / (2. * np.cos(2. * np.pi * Freq(np.arange(r0), s0) / s0) + 4.)
w1 = 6. / (2. * np.cos(2. * np.pi * Freq(np.arange(r1), s1) / s1) + 4.)
# FIXME: switch to pyfftw :
bic_filter = lambda _map: np.fft.irfft2(
np.fft.rfft2(_map) * np.outer(w0, w1))
dx = bic_filter(dx)
dy = bic_filter(dy)
Minv_xy = bic_filter(Minv_xy)
Minv_yx = bic_filter(Minv_yx)
Minv_xx = bic_filter(Minv_xx)
Minv_yy = bic_filter(Minv_yy)
header = r' "%s/lensit/gpu/bicubicspline.h" ' % os.path.abspath(os.curdir)
iterate = r"\
double fx,fy;\
double ex_len_dx,ey_len_dy,len_Mxx,len_Mxy,len_Myx,len_Myy;\
int i = 0;\
for(int y= 0; y < width; y++ )\
{\
for(int x = 0; x < width; x++,i++)\
{\
fx = x + ex[i];\
fy = y + ey[i];\
ex_len_dx = ex[i] + bicubiclensKernel(dx,fx,fy,width);\
ey_len_dy = ey[i] + bicubiclensKernel(dy,fx,fy,width);\
len_Mxx = bicubiclensKernel(Minv_xx,fx,fy,width);\
len_Myy = bicubiclensKernel(Minv_yy,fx,fy,width);\
len_Mxy = bicubiclensKernel(Minv_xy,fx,fy,width);\
len_Myx = bicubiclensKernel(Minv_yx,fx,fy,width);\
ex[i] += len_Mxx * ex_len_dx + len_Mxy * ey_len_dy;\
ey[i] += len_Myx * ex_len_dx + len_Myy * ey_len_dy;\
}\
}\
"
width = int(s0)
assert s0 == s1, 'Havent checked how this works with rectangular maps'
for i in range(0, NR_iter):
weave.inline(iterate, [
'ex', 'ey', 'dx', 'dy', 'Minv_xx', 'Minv_yy', 'Minv_xy', 'Minv_yx',
'width'
],
headers=[header])
return ex * rmin1, ey * rmin0
def get_inverse(self, NR_iter=None, use_Pool=0, crude=0, HD_res=None):
"""
:param NR_iter:
:param use_Pool: if positive, use Python multiprocessing Pool packacge.
if 0, serial calculation
if negative, send it on the GPU.
:param crude: Uses some crude scheme
:param HD_res: augmente the resolution to perform the inversion.
:return:
"""
if HD_res is not None:
# FIXME : this upgrade can be done in GPU for use_Pool < 0
lib_dir = self.lib_dir if self.lib_dir is None else self.lib_dir + '/temp_fup'
f_up = ffs_displacement(rfft2_utils.upgrade_map(self.get_dx(), HD_res),
rfft2_utils.upgrade_map(self.get_dy(), HD_res), self.lsides,
LD_res=self.LD_res, verbose=self.verbose, spline_order=self.k,
rule_for_derivative=self.rule,
NR_iter=self.NR_iter, lib_dir=lib_dir)
f_up_inv = f_up.get_inverse(NR_iter=NR_iter, use_Pool=use_Pool, crude=crude)
LD_res = Log2ofPowerof2(self.shape)
return ffs_displacement(rfft2_utils.subsample(f_up_inv.get_dx(), LD_res),
rfft2_utils.subsample(f_up_inv.get_dy(), LD_res), self.lsides,
LD_res=self.LD_res, verbose=self.verbose, spline_order=self.k,
rule_for_derivative=self.rule,
NR_iter=self.NR_iter, lib_dir=self.lib_dir)
if crude > 0:
return self.get_inverse_crude(crude)
if NR_iter is None: NR_iter = self.NR_iter
if use_Pool > 0:
if not self.is_dxdy_ondisk():
assert self.has_lib_dir(), "Specify lib. dir. if you want to use Pool."
print "lens_map::writing displacements on disk :"
if not os.path.exists(self.lib_dir): os.makedirs(self.lib_dir)
self.write_npy(self.lib_dir + '/temp_displ' + str(pbs.rank)) # this turns dx and dy to the paths
dx_inv, dy_inv = ffs_pool.get_inverse_Pooled(self.mk_args('', self.dx, self.dy, NR_iter=NR_iter),
root_Nthreads=abs(use_Pool))
return ffs_displacement(dx_inv, dy_inv, self.lsides, lib_dir=self.lib_dir,
LD_res=self.LD_res, verbose=self.verbose, spline_order=self.k, NR_iter=self.NR_iter)
elif use_Pool == 0:
spliter_lib = map_spliter.periodicmap_spliter() # library to split periodic maps.
dx_inv, dy_inv = np.empty(self.shape), np.empty(self.shape)
label = 'ffs_deflect::calculating inverse displ. field'
for i, N in utils.enumerate_progress(xrange(self.N_chks), label=label):
# Doing chunk N
dx_inv_N, dy_inv_N = self.get_inverse_chk_N(N, NR_iter=NR_iter)
sLDs, sHDs = spliter_lib.get_slices_chk_N(N, self.LD_res, self.HD_res, self.buffers, inverse=True)
# Pasting it onto the full map
dx_inv[sHDs[0]] = dx_inv_N[sLDs[0]]
dy_inv[sHDs[0]] = dy_inv_N[sLDs[0]]
return ffs_displacement(dx_inv, dy_inv, self.lsides, lib_dir=self.lib_dir,
LD_res=self.LD_res, verbose=self.verbose, spline_order=self.k, NR_iter=self.NR_iter)
elif use_Pool < 0:
# GPU calculation.
from lensit.gpu import inverse_GPU as inverse_GPU
GPU_res = np.array(inverse_GPU.GPU_HDres_max)
if np.all(np.array(self.HD_res) <= GPU_res):
# No need to split maps :
dx_inv, dy_inv = inverse_GPU.inverse_GPU(self.get_dx(), self.get_dy(), self.rmin, NR_iter)
return ffs_displacement(dx_inv, dy_inv, self.lsides, lib_dir=self.lib_dir,
LD_res=self.LD_res, verbose=self.verbose, spline_order=self.k,
NR_iter=self.NR_iter)
else:
LD_res, buffers = get_GPUbuffers(GPU_res)
assert np.all(np.array(buffers) > (np.array(self.buffers) + 5.)), (buffers, self.buffers)
Nchunks = 2 ** (np.sum(np.array(self.HD_res) - np.array(LD_res)))
dx_N = np.empty((2 ** LD_res[0] + 2 * buffers[0], 2 ** LD_res[1] + 2 * buffers[1]))
dy_N = np.empty((2 ** LD_res[0] + 2 * buffers[0], 2 ** LD_res[1] + 2 * buffers[1]))
if self.verbose:
print '++ inverse displacement :' \
' splitting inverse on GPU , chunk shape %s, buffers %s' % (dx_N.shape, buffers)
spliter_lib = map_spliter.periodicmap_spliter() # library to split periodic maps.
dx_inv, dy_inv = np.empty(self.shape), np.empty(self.shape) # Outputs
for N in xrange(Nchunks):
sLDs, sHDs = spliter_lib.get_slices_chk_N(N, LD_res, self.HD_res, buffers)
for sLD, sHD in zip(sLDs, sHDs):
dx_N[sLD] = self.get_dx()[sHD]
dy_N[sLD] = self.get_dy()[sHD]
dx_inv_N, dy_inv_N = inverse_GPU.inverse_GPU(dx_N, dy_N, self.rmin, NR_iter)
sLDs, sHDs = spliter_lib.get_slices_chk_N(N, LD_res, self.HD_res, buffers, inverse=True)
dx_inv[sHDs[0]] = dx_inv_N[sLDs[0]]
dy_inv[sHDs[0]] = dy_inv_N[sLDs[0]]
return ffs_displacement(dx_inv, dy_inv, self.lsides, lib_dir=self.lib_dir,
LD_res=self.LD_res, verbose=self.verbose, spline_order=self.k,
NR_iter=self.NR_iter)
elif use_Pool == 100:
assert 0
else:
assert 0
def lens_map(self, map, use_Pool=0, tidy=True, crude=0, do_not_prefilter=False):
"""
Lens the input map according to the displacement fields dx dy. 'map' typically could be (8192 * 8192) np array,
or the path to the array on disk.
Does this by splitting the job in chunks (of typically (256 * 256), as specified by the LD_res parameters)
allowing a buffer size to ensure the junctions are properly performed.
Set use_Pool to a power of two to use explicit threading via the multiprocessing module, or, if < 0,
to perform the operation on the GPU.
if > 0 'use_Pool' ** 2 is the number of threads. On laptop and Darwin use_Pool = 16 has the best performances.
It use_Pool is set, then 'map' must be the path to the map to lens or map will be saved to disk.
"""
# TODO : could evaluate the splines at low res.
assert self.load_map(map).shape == self.shape, (self.load_map(map).shape, self.shape)
if crude > 0:
return self.lens_map_crude(map, crude)
if use_Pool < 0:
# use of GPU :
try:
from lensit.gpu import lens_GPU
except:
assert 0, 'Import of mllens lens_GPU failed !'
GPU_res = np.array(lens_GPU.GPU_HDres_max)
if np.all(np.array(self.HD_res) <= GPU_res):
return lens_GPU.lens_onGPU(map, self.get_dx_ingridunits(), self.get_dy_ingridunits(),
do_not_prefilter=do_not_prefilter)
LD_res, buffers = get_GPUbuffers(GPU_res)
assert np.all(np.array(buffers) > (np.array(self.buffers) + 5.)), (buffers, self.buffers)
Nchunks = 2 ** (np.sum(np.array(self.HD_res) - np.array(LD_res)))
lensed_map = np.empty(self.shape) # Output
dx_N = np.empty((2 ** LD_res[0] + 2 * buffers[0], 2 ** LD_res[1] + 2 * buffers[1]))
dy_N = np.empty((2 ** LD_res[0] + 2 * buffers[0], 2 ** LD_res[1] + 2 * buffers[1]))
unl_CMBN = np.empty((2 ** LD_res[0] + 2 * buffers[0], 2 ** LD_res[1] + 2 * buffers[1]))
if self.verbose:
print '++ lensing map :' \
' splitting map on GPU , chunk shape %s, buffers %s' % (dx_N.shape, buffers)
spliter_lib = map_spliter.periodicmap_spliter() # library to split periodic maps.
for N in xrange(Nchunks):
sLDs, sHDs = spliter_lib.get_slices_chk_N(N, LD_res, self.HD_res, buffers)
for sLD, sHD in zip(sLDs, sHDs):
dx_N[sLD] = self.get_dx()[sHD] / self.rmin[1]
dy_N[sLD] = self.get_dy()[sHD] / self.rmin[0]
unl_CMBN[sLD] = self.load_map(map)[sHD]
sLDs, sHDs = spliter_lib.get_slices_chk_N(N, LD_res, self.HD_res, buffers, inverse=True)
lensed_map[sHDs[0]] = lens_GPU.lens_onGPU(unl_CMBN, dx_N, dy_N, do_not_prefilter=do_not_prefilter)[
sLDs[0]]
return lensed_map
elif use_Pool > 100:
if not isinstance(map, str):
assert self.has_lib_dir(), "Specify lib. dir. if you want to use Pool."
np.save(self.lib_dir + '/temp_maptolens_rank%s.npy' % pbs.rank, map)
if not self.is_dxdy_ondisk():
assert self.has_lib_dir(), "Specify lib. dir. if you want to use Pool."
print "lens_map::writing displacements on disk :"
self.write_npy(self.lib_dir + '/temp_displ' + str(pbs.rank)) # this turns dx and dy to the paths
path_to_map = map if isinstance(map, str) else self.lib_dir + '/temp_maptolens_rank%s.npy' % pbs.rank
ret = ffs_pool.get_lens_Pooled(self.mk_args(path_to_map, self.dx, self.dy),
root_Nthreads=use_Pool % 100, do_not_prefilter=do_not_prefilter)
if tidy:
if os.path.exists(self.lib_dir + '/temp_maptolens_rank%s.npy' % pbs.rank):
os.remove(self.lib_dir + '/temp_maptolens_rank%s.npy' % pbs.rank)
return ret
elif use_Pool == 100 or use_Pool == 101:
assert self.load_map(map).shape == self.shape, self.load_map(map).shape
s = self.chk_shape
idc0, idc1 = np.indices(s) # Two (256 * 256) maps
dx_gu = np.empty(s) # will dx displ. in grid units of each chunk (typ. (256 * 256) )
dy_gu = np.empty(s) # will dy displ. in grid units of each chunk (typ. (256 * 256) )
map_chk = np.empty(s) # Will be map chunk
# (typ. (256 * 256) )
lensed_map = np.empty(self.shape)
spliter_lib = map_spliter.periodicmap_spliter() # library to split periodic maps.
for N in xrange(self.N_chks):
# doing chunk N
sLDs, sHDs = spliter_lib.get_slices_chk_N(N, self.LD_res, self.HD_res, self.buffers)
for sLD, sHD in zip(sLDs, sHDs):
# Displacements chunk in grid units, and map chunk to displace.
dx_gu[sLD] = self.get_dx()[sHD] / self.rmin[1]
dy_gu[sLD] = self.get_dy()[sHD] / self.rmin[0]
map_chk[sLD] = self.load_map(map)[sHD]
if do_not_prefilter:
# Undoing the prefiltering prior to apply bicubic interpolation
map_chk = np.fft.rfft2(map_chk)
w0 = 6. / (2. * np.cos(2. * np.pi * np.fft.fftfreq(s[0])) + 4.)
map_chk /= np.outer(w0, w0[0:map_chk.shape[1]])
map_chk = np.fft.irfft2(map_chk, s)
lx = (idc1 + dx_gu).flatten() # No need to enforce periodicity here.
ly = (idc0 + dy_gu).flatten() # No need to enforce periodicity here.
sLDs, sHDs = spliter_lib.get_slices_chk_N(N, self.LD_res, self.HD_res, self.buffers, inverse=True)
lensed_map[sHDs[0]] = interpolate.RectBivariateSpline(np.arange(s[0]), np.arange(s[1]),
map_chk, kx=self.k, ky=self.k).ev(ly, lx).reshape(
self.chk_shape)[sLDs[0]]
return lensed_map
elif use_Pool == 0 or use_Pool == 1:
assert self.shape[0] == self.shape[1], self.shape
bicubicspline = r"\
int i,j;\
for( j= 0; j < width; j++ )\
{\
for( i = 0; i < width; i++)\
{\
lenmap[j * width + i] = bicubiclensKernel(filtmap,i + dx_gu[j * width + i],j + dy_gu[j * width + i],width);\
}\
}"
header = r' "%s/lensit/gpu/bicubicspline.h" ' % li.LENSITDIR
if do_not_prefilter:
filtmap = self.load_map(map).astype(np.float64)
else:
# TODO : may want to add pyFFTW here as well
filtmap = np.fft.rfft2(self.load_map(map))
w0 = 6. / (2. * np.cos(2. * np.pi * np.fft.fftfreq(filtmap.shape[0])) + 4.)
filtmap *= np.outer(w0, w0[0:filtmap.shape[1]])
filtmap = np.fft.irfft2(filtmap, self.shape)
lenmap = np.empty(self.shape, dtype=np.float64)
dx_gu = self.get_dx_ingridunits().astype(np.float64)
dy_gu = self.get_dy_ingridunits().astype(np.float64)
width = int(self.shape[0])
assert self.shape[0] == self.shape[1]
weave.inline(bicubicspline, ['lenmap', 'filtmap', 'dx_gu', 'dy_gu', 'width'], headers=[header])
return lenmap
elif use_Pool > 1 and use_Pool < 100:
# TODO : may want to add pyFFTW here as well
if not isinstance(map, str):
assert self.has_lib_dir(), "Specify lib. dir. if you want to use Pool."
np.save(self.lib_dir + '/temp_maptolens_rank%s.npy' % pbs.rank, map)
if not self.is_dxdy_ondisk():
assert self.has_lib_dir(), "Specify lib. dir. if you want to use Pool."
print "lens_map::writing displacements on disk :"
self.write_npy(self.lib_dir + '/temp_displ' + str(pbs.rank)) # this turns dx and dy to the paths
path_to_map = map if isinstance(map, str) else self.lib_dir + '/temp_maptolens_rank%s.npy' % pbs.rank
ret = ffs_pool.get_lens_Pooled_weave(self.mk_args(path_to_map, self.dx, self.dy),
root_Nthreads=use_Pool, do_not_prefilter=do_not_prefilter)
if tidy:
if os.path.exists(self.lib_dir + '/temp_maptolens_rank%s.npy' % pbs.rank):
os.remove(self.lib_dir + '/temp_maptolens_rank%s.npy' % pbs.rank)
return ret
else:
assert 0
