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)