def get_inverse(self,NR_iter= None,use_Pool = 0):
"""
Returns inverse displacement of the full map.
"""
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 :"
self.write_npy( self.lib_dir + '/temp_displ' + str(pbs.rank)) # this turns dx and dy to the paths
dx_inv,dy_inv = lens_Pool.get_inverse_Pooled(self.mk_args('',self.dx,self.dy),root_Nthreads = use_Pool)
return ffs_displacement(dx_inv,dy_inv,lib_dir=self.lib_dir,
LD_res=self.LD_res,verbose = self.verbose,spline_order=self.k,NR_iter= self.NR_iter)
if NR_iter is None : NR_iter = self.NR_iter
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_displacement::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,lib_dir=self.lib_dir,
LD_res=self.LD_res,verbose = self.verbose,spline_order=self.k,NR_iter= self.NR_iter)
def _lens_chk_N(args):
assert len(args) == 12,args
N,path_to_map,path_to_dx,path_to_dy,buff0,buff1,HD_res0,HD_res1,NR_iter,kspl,LD0,LD1 = 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')
# TODO : pass this as argument ?
rmin0 = np.sqrt(4 * np.pi)/ 2 ** HD_res[0]
rmin1 = np.sqrt(4 * np.pi)/ 2 ** HD_res[1]
map_chk_N = np.zeros(s)
dx_gu = np.zeros(s) # will dx displ. in grid units of each chunk (typ. (256 * 256) )
dy_gu = np.zeros(s) # 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]
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 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_map(self,map,use_Pool = 0):
"""
Lens the input map according to the displacement fields dx dy. 'map' typically would 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.
'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.
"""
if use_Pool > 0 :
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.npy',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.npy'
return lens_Pool.get_lens_Pooled(self.mk_args(path_to_map,self.dx,self.dy),root_Nthreads = use_Pool)
assert map.shape == self.shape,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 i,N in utils.enumerate_progress(xrange(self.N_chks),label = 'ffs_displacement::lensing map'):
# 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] = map[sHD]
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
def lens_map(self, map, split=None, verbose=None):
" Returns out_map(y,x) = in_map(y+ dy,x + dx) using splines by default "
assert np.all(map.shape == self.shape), "Input not understood" + str(map.shape) + ' / ' + str(self.shape)
if verbose is None: verbose = self.verbose
if split is None: split = self.split
rmin = self.rmin()
if split:
# we split the map in 4 to avoid boundary effects, due to the non period. of RectBivariateSpline.
spliter = map_spliter.periodicmap_spliter()
buffer0 = 6 * np.int64(np.max(np.abs(self.dy)) / rmin[0])
buffer1 = 6 * np.int64(np.max(np.abs(self.dx)) / rmin[1])
buffers = (buffer0, buffer1)
if verbose: print 'lens_map::split buffers in lens_map', buffers
len_chks = []
for i in xrange(4):
# It is also possible to call all the chunks at once.
# In the current implementation of split this version consumes less memory but is less efficient.
map_chk = spliter.get_splitin4_chunk_i(map, buffers, i)
dx_chk = spliter.get_splitin4_chunk_i(self.dx, buffers, i)
dy_chk = spliter.get_splitin4_chunk_i(self.dy, buffers, i)
s = map_chk.shape
idc = np.indices(s)
lx = (idc[1] * rmin[1] + dx_chk).flatten()
ly = (idc[0] * rmin[0] + dy_chk).flatten()
# No periodicty.By buffer constr. the points should still be in the map.
len_chks.append(interpolate.RectBivariateSpline(np.arange(s[0]) * rmin[0], np.arange(s[1]) * rmin[1],
map_chk, kx=self.k, ky=self.k).ev(ly, lx).reshape(s))
return spliter.splitin4_inverse(len_chks, buffers)
if verbose: print 'no split in lens_map, you may lack of accuracy at the boundaries'
s = self.shape
idc = np.indices(s)
# If we do not split at least we take into account periodicity of the coordinates.
lx = (idc[1] * rmin[1] + self.dx).flatten() % self.lsides[1]
ly = (idc[0] * rmin[0] + self.dy).flatten() % self.lsides[0]
newmap = interpolate.RectBivariateSpline(np.arange(s[0]) * rmin[0], np.arange(s[1]) * rmin[1], map, kx=self.k,
ky=self.k).ev(ly, lx).reshape(s)
return newmap
def Pool_generic(func,arg,root_Nthreads):
assert len(arg) == 9,arg
path_to_map,path_to_dx,path_to_dy,buff0,buff1,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,HD_res0,HD_res1,NR_iter,kspl,LD[0],LD[1]]
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)
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)
return ret
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) == 12,args
N,path_to_map,path_to_dx,path_to_dy,buff0,buff1,HD_res0,HD_res1,NR_iter,kspl,LD0,LD1 = args
HD_res = (HD_res0,HD_res1)
LD = (LD0,LD1)
s = (2 ** LD[0] + 2 * buff0,2 ** LD[1] + 2 * buff1) # Chunk shape
rmin0 = np.sqrt(4 * np.pi)/ 2 ** HD_res[0]
rmin1 = np.sqrt(4 * np.pi)/ 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.
rmin0 = np.sqrt(4 * np.pi)/ 2 ** HD_res[0]
rmin1 = np.sqrt(4 * np.pi)/ 2 ** HD_res[1]
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 inverse(self, NR_iter=None, split=None):
"""
Returns the inverse displacement field through Newton-Raphson iterations.
Attainable accuracy depends on resolution.
With PL2015 cell size and parameters, NR_iter = 1 is near-optimal already and achieve
0.1 uK accuracy on lensing - delensing test. See the ipython notebook on displacements test.
Make sure to use the 'split' option in lens_map and this routine, in order to handle properly boundaries.
If not the error is greater by a factor of 10. (still, not too bad depending on what you want I guess)
"""
verbose = self.verbose
if split is None: split = self.split
if NR_iter is None: NR_iter = self.NR_iter
s = self.shape
rmin = self.rmin()
det = self.det_magn()
# Inverse magn. elements. (with a minus sign) We may need to spline these later for further NR iterations
_Minv_xx = -(1. + self.M_mat['kappa'] - self.M_mat['gamma_Q']) / det
_Minv_yy = -(1. + self.M_mat['kappa'] + self.M_mat['gamma_Q']) / det
_Minv_xy = (self.M_mat['gamma_U'] + self.M_mat['omega']) / det
_Minv_yx = (self.M_mat['gamma_U'] - self.M_mat['omega']) / det
# 0th NR step.
dxn = (_Minv_xx * self.dx + _Minv_xy * self.dy)
dyn = (_Minv_yy * self.dy + _Minv_yx * self.dx)
if NR_iter == 0: return displacement_field_flatsky_physical(dxn, dyn, self.lsides, verbose=self.verbose,
spline_order=self.k, rule_for_derivative=self.rule,
initiate_magn=False)
# Setting up a bunch of splines to interpolate the increment to the displacement according to NR.
# 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.
# FIXME : Think on how to best split the NR inversion. The 0th step is OK (SURE ?), no need to worry.
# FIXME : Recursive implementation would then be cleaner. But it looks like it's working as such.
if split:
spliter = map_spliter.periodicmap_spliter()
buffer0 = 6 * np.int64(np.max(np.abs(self.dy)) / rmin[0])
buffer1 = 6 * np.int64(np.max(np.abs(self.dx)) / rmin[1])
buffers = (buffer0, buffer1)
if verbose: print 'inverse_displacement:: split buffers ', buffers
dx_sol_chks = []
dy_sol_chks = []
for isplit in xrange(4):
dx_chk = spliter.get_splitin4_chunk_i(self.dx, buffers, isplit)
dy_chk = spliter.get_splitin4_chunk_i(self.dy, buffers, isplit)
_Minv_xx_chk = spliter.get_splitin4_chunk_i(_Minv_xx, buffers, isplit)
_Minv_xy_chk = spliter.get_splitin4_chunk_i(_Minv_xy, buffers, isplit)
_Minv_yy_chk = spliter.get_splitin4_chunk_i(_Minv_yy, buffers, isplit)
_Minv_yx_chk = spliter.get_splitin4_chunk_i(_Minv_yx, buffers, isplit)
s = dx_chk.shape
xcoord = np.arange(s[1]) * rmin[1]
ycoord = np.arange(s[0]) * rmin[0]
spl_dx = interpolate.RectBivariateSpline(ycoord, xcoord, dx_chk, kx=self.k, ky=self.k)
spl_dy = interpolate.RectBivariateSpline(ycoord, xcoord, dy_chk, kx=self.k, ky=self.k)
spl_xx = interpolate.RectBivariateSpline(ycoord, xcoord, _Minv_xx_chk, kx=self.k, ky=self.k)
spl_yy = interpolate.RectBivariateSpline(ycoord, xcoord, _Minv_yy_chk, kx=self.k, ky=self.k)
spl_xy = interpolate.RectBivariateSpline(ycoord, xcoord, _Minv_xy_chk, kx=self.k, ky=self.k)
spl_yx = interpolate.RectBivariateSpline(ycoord, xcoord, _Minv_yx_chk, kx=self.k, ky=self.k)
idc = np.indices(s)
y_x = idc[1] * rmin[1]
y_y = idc[0] * rmin[0]
del idc
for i in range(NR_iter):
if self.verbose: print "Additional NR step", i + 1, "split # : ", isplit
# If this is the first call we need to get the starting point from the unsplited 0th solution :
dxn_1 = spliter.get_splitin4_chunk_i(dxn, buffers, isplit) if i == 0 else dxn_i
dyn_1 = spliter.get_splitin4_chunk_i(dyn, buffers, isplit) if i == 0 else dyn_i
# No need to consider periodicity in the chunks by construction
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_i = dxn_1 + spl_xx.ev(ly, lx).reshape(s) * res_x + spl_xy.ev(ly, lx).reshape(s) * res_y
dyn_i = dyn_1 + spl_yx.ev(ly, lx).reshape(s) * res_x + spl_yy.ev(ly, lx).reshape(s) * res_y
dx_sol_chks.append(dxn_i)
dy_sol_chks.append(dyn_i)
# Recombine solutions and returns :
dxn = spliter.splitin4_inverse(dx_sol_chks, buffers)
dyn = spliter.splitin4_inverse(dy_sol_chks, buffers)
return displacement_field_flatsky_physical(dxn, dyn, self.lsides, initiate_magn=False,
verbose=self.verbose, spline_order=self.k,
rule_for_derivative=self.rule)
# No splitting :
if verbose: print 'inverse_displacement :: no split in NR inversion, you may get boudary effects'
xcoord = np.arange(s[1]) * rmin[1]
ycoord = np.arange(s[0]) * rmin[0]
spl_dx = interpolate.RectBivariateSpline(ycoord, xcoord, self.dx, kx=self.k, ky=self.k)
spl_dy = interpolate.RectBivariateSpline(ycoord, xcoord, self.dy, kx=self.k, ky=self.k)
spl_xx = interpolate.RectBivariateSpline(ycoord, xcoord, _Minv_xx, kx=self.k, ky=self.k)
spl_yy = interpolate.RectBivariateSpline(ycoord, xcoord, _Minv_yy, kx=self.k, ky=self.k)
spl_xy = interpolate.RectBivariateSpline(ycoord, xcoord, _Minv_xy, kx=self.k, ky=self.k)
spl_yx = interpolate.RectBivariateSpline(ycoord, xcoord, _Minv_yx, kx=self.k, ky=self.k)
idc = np.indices(s)
y_x = idc[1] * rmin[1]
y_y = idc[0] * rmin[0]
del idc
for i in range(NR_iter):
if self.verbose: print "Additional NR step : ", i + 1
dxn_1 = dxn
dyn_1 = dyn
lx = (y_x + dxn_1).flatten() % self.lsides[1] # % due to periodicity of the box
ly = (y_y + dyn_1).flatten() % self.lsides[0]
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 displacement_field_flatsky_physical(dxn, dyn, self.lsides, initiate_magn=False,
verbose=self.verbose, spline_order=self.k,
rule_for_derivative=self.rule)
