import numpy as np """ Numpy performs by default the 1/N scaling with the backward transform (ifft) and not the forward as shown in (2). """ ########################################################## ############### FFT Convolution Routines ################# ########################################################## def CONVOLVE_3D(signal_1,signal_2): signal_conv = np.fft.fftshift(np.fft.ifftn(np.fft.fftn(signal_1) * np.fft.fftn(signal_2))) return signal_conv def COMPUTE_FFT_3D_NO_NORM_NO_SHIFT(signal): # No normalization and no shifting is required for convolution operation signal_fft = np.zeros(signal.shape,dtype=np.complex64) signal_fft[:,:,:] = np.fft.fftn(signal) return signal_fft def CONVOLVE_3D_WITH_FFT_INPUT(signal_1_fft,signal_2_fft): signal_conv = np.zeros(signal_1_fft.shape,dtype=np.float32) signal_conv[:,:,:] = np.abs(np.fft.fftshift(np.fft.ifftn(signal_1_fft * signal_2_fft))) return signal_conv def COMPUTE_DERIVATIVE_USING_FFT(signal_arr_3d,k_arr_3d): # Possible to enhance the speed of the code here. if np.all(signal_arr_3d) == 0.0: signal_derivative = signal_arr_3d else: signal_fft_arr = np.fft.fftshift(np.fft.fftn(signal_arr_3d)) signal_derivative = np.real(np.fft.ifftn(np.fft.ifftshift(1j*2.0*np.pi*k_arr_3d*signal_fft_arr))) return signal_derivative def COMPUTE_DERIVATIVE_USING_GRADIENT(signal_arr_3d,axis_indx): # Possible to enhance the speed of the code here. signal_derivative = np.gradient(signal_arr_3d,axis=axis_indx) return signal_derivative def BACKUP_COMPUTE_DERIVATIVE_USING_GRADIENT(signal_arr_3d,inp_dict,axis_indx): # Possible to enhance the speed of the code here. N_Grid = inp_dict["N_Grid"] cosmo_param = inp_dict["cosmo_param"] hubble_param = cosmo_param["hubble_param"] astro_const = inp_dict["astro_const"] mpc_to_cm = astro_const["mpc_to_cm"] #dx = inp_dict["L_box"] * mpc_to_cm / hubble_param /N_Grid / 16.0 #dx = inp_dict["L_box"] / N_Grid signal_derivative = np.gradient(signal_arr_3d,axis=axis_indx) return signal_derivative ########################################################## ####### Real Space array (1D and 3D) calculations ######## ########################################################## def COMPUTE_1D_FOURIER_SPACE_ARRAY_IN_BOX_SIZE_UNITS(inp_dict): """ Compute 3D Cartesian box size This routine computes 3D grid in integers like (-dx*N_Grid/2 , ... , -dx , 0 , dx , ... , dx*N_Grid/2) """ key_list_reqd = ["L_box","N_Grid"] CHECK_IF_KEY_PRESENT(key_list_reqd,inp_dict) #cosmo_param = inp_dict["cosmo_param"] #hubble_param = cosmo_param["hubble_param"] #astro_const = inp_dict["astro_const"] #mpc_to_cm = astro_const["mpc_to_cm"] #L_box = inp_dict["L_box"] * mpc_to_cm / hubble_param / 16.0 L_box = inp_dict["L_box"] L_box_by_2 = L_box * 0.5 N_Grid = inp_dict["N_Grid"] x_arr_1d = np.linspace(-L_box_by_2,L_box_by_2,N_Grid) dx = np.median(x_arr_1d[1:] - x_arr_1d[:-1]) k_arr_1d = np.fft.fftshift(np.fft.fftfreq(N_Grid,dx)) # Note shifting is applied inp_dict.update({"k_arr_1d":k_arr_1d}) return inp_dict def COMPUTE_3D_FOURIER_SPACE_ARRAY_IN_BOX_SIZE_UNITS(inp_dict,flag_norm="n"): """ Compute 3D Cartesian box size This routine computes 3D grid in integers like (-dx*N_Grid/2 , ... , -dx , 0 , dx , ... , dx*N_Grid/2) if flag_norm is "y" then k_arr_3d = (kx^2 + ky^2 + kz^2)^0.5 Else k_arr_3d would be 4d array """ if "k_arr_1d" not in inp_dict.keys(): inp_dict = COMPUTE_1D_FOURIER_SPACE_ARRAY_IN_BOX_SIZE_UNITS(inp_dict) k_arr_1d = inp_dict["k_arr_1d"] k_arr_3d_tup = np.meshgrid(k_arr_1d,k_arr_1d,k_arr_1d,indexing = 'xy') N_Grid = k_arr_3d_tup[0].shape[0] if flag_norm.lower() == "n": k_arr_3d = np.zeros((N_Grid,N_Grid,N_Grid,3),dtype = np.float32) k_arr_3d[:,:,:,0] = k_arr_3d_tup[0] k_arr_3d[:,:,:,1] = k_arr_3d_tup[1] k_arr_3d[:,:,:,2] = k_arr_3d_tup[2] elif flag_norm.lower() == "y": k_arr_3d = np.zeros((N_Grid,N_Grid,N_Grid),dtype = np.float32) k_arr_3d[:,:,:] += k_arr_3d_tup[0]**2.0 k_arr_3d[:,:,:] += k_arr_3d_tup[1]**2.0 k_arr_3d[:,:,:] += k_arr_3d_tup[2]**2.0 k_arr_3d[:,:,:] = k_arr_3d**0.5 return k_arr_3d ########################################################## ####### Real Space array (1D and 3D) calculations ######## ########################################################## def COMPUTE_1D_REAL_SPACE_ARRAY_IN_BOX_SIZE_UNITS(inp_dict): """ Compute 3D Cartesian box size This routine computes 3D grid in integers like (-dx*N_Grid/2 , ... , -dx , 0 , dx , ... , dx*N_Grid/2) """ key_list_reqd = ["L_box","N_Grid"] CHECK_IF_KEY_PRESENT(key_list_reqd,inp_dict) L_box = inp_dict["L_box"] L_box_by_2 = L_box * 0.5 N_Grid = inp_dict["N_Grid"] x_arr_1d = np.linspace(-L_box_by_2,L_box_by_2,N_Grid) inp_dict.update({"x_arr_1d":x_arr_1d}) return inp_dict def COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_SIZE_UNITS(inp_dict): """ Compute 3D Cartesian box size This routine computes 3D grid in integers like (-dx*N_Grid/2 , ... , -dx , 0 , dx , ... , dx*N_Grid/2) """ if "x_arr_1d" not in inp_dict.keys(): inp_dict = COMPUTE_1D_REAL_SPACE_ARRAY_IN_BOX_SIZE_UNITS(inp_dict) x_arr_1d = inp_dict["x_arr_1d"] x_arr_3d_tup = np.meshgrid(x_arr_1d,x_arr_1d,x_arr_1d,indexing = 'xy') x_arr_3d = x_arr_3d_tup[0] y_arr_3d = x_arr_3d_tup[1] z_arr_3d = x_arr_3d_tup[2] x_mod_arr_3d = (x_arr_3d**2.0 + y_arr_3d**2.0 + z_arr_3d**2.0)**0.5 return x_mod_arr_3d ########################################################## ####### Real Space array (1D and 3D) calculations ######## ########################################################## def COMPUTE_1D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict): """ Compute 3D Cartesian Grid This routine computes 3D grid in integers like (-512,-511, ... ,0, ... ,510,511) """ key_list_reqd = ["L_box","N_Grid"] CHECK_IF_KEY_PRESENT(key_list_reqd,inp_dict) L_box = inp_dict["L_box"] L_box_by_2 = L_box * 0.5 N_Grid = inp_dict["N_Grid"] x_arr_1d = np.linspace(-N_Grid/2,N_Grid/2-1,N_Grid) inp_dict.update({"x_arr_1d":x_arr_1d}) return inp_dict def COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict,flag_return_all="n"): """ Compute 3D Cartesian Grid This routine computes 3D grid in integers like (-512,-511, ... ,0, ... ,510,511) """ if "x_arr_1d" not in inp_dict.keys(): inp_dict = COMPUTE_1D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict) x_arr_1d = inp_dict["x_arr_1d"] x_arr_3d_tup = np.meshgrid(x_arr_1d,x_arr_1d,x_arr_1d,indexing = 'xy') x_arr_3d = x_arr_3d_tup[0] y_arr_3d = x_arr_3d_tup[1] z_arr_3d = x_arr_3d_tup[2] x_mod_arr_3d = (x_arr_3d**2.0 + y_arr_3d**2.0 + z_arr_3d**2.0)**0.5 if flag_return_all == "n": return x_mod_arr_3d elif flag_return_all == "y": return x_arr_3d,y_arr_3d,z_arr_3d,x_mod_arr_3d def COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX_POLAR_COORDINATES(inp_dict): """ Compute 3D Cartesian Grid This routine computes 3D grid in integers like (-512,-511, ... ,0, ... ,510,511) """ if "x_arr_1d" not in inp_dict.keys(): inp_dict = COMPUTE_1D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict) x_arr_1d = inp_dict["x_arr_1d"] x_arr_3d_tup = np.meshgrid(x_arr_1d,x_arr_1d,x_arr_1d,indexing = 'xy') x_arr_3d = x_arr_3d_tup[0] y_arr_3d = x_arr_3d_tup[1] z_arr_3d = x_arr_3d_tup[2] xy_arr_3d_sq = x_arr_3d**2.0 + y_arr_3d**2.0 R_arr_3d = (xy_arr_3d_sq + z_arr_3d**2.0)**0.5 theta_arr_3d = np.arctan2(xy_arr_3d_sq**0.5, z_arr_3d) # for elevation angle defined from Z-axis down #theta_arr_3d = np.arctan2(z_arr_3d,xy_arr_3d_sq**0.5) # for elevation angle defined from XY-plane up phi_arr_3d = np.arctan2(x_arr_3d, y_arr_3d) return R_arr_3d,theta_arr_3d,phi_arr_3d ########################################################## ####### Compute Window function in Fourier Space ######### ########################################################## def COMPUTE_WINDOW_FUNCTION_REAL_SPACE(inp_dict,flag_window="top_hat"): if flag_window == "top_hat": # Compute Spherical Top Hat Function WR_arr_3d,npix = COMPUTE_SPHERICAL_TOP_HAT_REAL_SPACE(inp_dict) return WR_arr_3d,npix elif flag_window == "top_hat_shell": # Compute Spherical Top Hat Function WR_arr_3d,npix = COMPUTE_SPHERICAL_SHELL_TOP_HAT_REAL_SPACE(inp_dict) return WR_arr_3d,npix elif flag_window == "gaussian": WR_arr_3d,npix = COMPUTE_GAUSSIAN_REAL_SPACE(inp_dict) return WR_arr_3d,npix elif flag_window == "exponential": WR_arr_3d,npix = COMPUTE_EXPONENTIAL_REAL_SPACE(inp_dict) return WR_arr_3d,npix elif flag_window == "top_hat_sectors": WR_arr_3d,npix,sector_info_arr = COMPUTE_SPHERICAL_TOP_HAT_WITH_SLICES_REAL_SPACE(inp_dict) return WR_arr_3d,npix,sector_info_arr elif flag_window == "top_hat_subcubes": WR_arr_3d,npix,sector_info_arr = COMPUTE_SPHERICAL_TOP_HAT_WITH_SUBCUBES_REAL_SPACE(inp_dict) return WR_arr_3d,npix,sector_info_arr def COMPUTE_SPHERICAL_TOP_HAT_REAL_SPACE(inp_dict): """ Very Important Spherical top hat is calculated in terms of grid size and not comoving or physical distances. This makes easier to use same code for different box sizes. """ N_Grid = inp_dict["N_Grid"] R_sphere = inp_dict["R_sphere"] x_mod_arr_3d = COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict) WR_arr_3d = np.zeros((N_Grid,N_Grid,N_Grid)) bool_arr = x_mod_arr_3d <= R_sphere npix = np.sum(bool_arr) V_sphere_inv = 1.0 / float(npix) # Normalization constant WR_arr_3d[bool_arr] = V_sphere_inv return WR_arr_3d,npix def COMPUTE_SPHERICAL_SHELL_TOP_HAT_REAL_SPACE(inp_dict): """ Very Important Spherical top hat is calculated in terms of grid size and not comoving or physical distances. This makes easier to use same code for different box sizes. """ N_Grid = inp_dict["N_Grid"] R_sphere_min = inp_dict["R_sphere_min"] R_sphere_max = inp_dict["R_sphere_max"] x_mod_arr_3d = COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict) WR_arr_3d = np.zeros((N_Grid,N_Grid,N_Grid)) bool_arr = (R_sphere_min <= x_mod_arr_3d) & (x_mod_arr_3d < R_sphere_max) npix = np.sum(bool_arr) if npix > 0: V_sphere_inv = 1.0 / float(npix) # Normalization constant WR_arr_3d[bool_arr] = V_sphere_inv return WR_arr_3d,npix def COMPUTE_GAUSSIAN_REAL_SPACE(inp_dict): """ Very Important Spherical top hat is calculated in terms of grid size and not comoving or physical distances. This makes easier to use same code for different box sizes. """ N_Grid = inp_dict["N_Grid"] R_sphere = inp_dict["R_sphere"] x_mod_arr_3d = COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict) WR_arr_3d = np.zeros((N_Grid,N_Grid,N_Grid)) WR_arr_3d[:,:,:] = np.exp(-0.5*x_mod_arr_3d**2.0/R_sphere**2.0) WR_sum = np.sum(WR_arr_3d) npix = 0.0 #np.sum(WR_arr_3d) WR_arr_3d /= WR_sum return WR_arr_3d,npix def COMPUTE_EXPONENTIAL_REAL_SPACE(inp_dict): """ Very Important Spherical top hat is calculated in terms of grid size and not comoving or physical distances. This makes easier to use same code for different box sizes. """ N_Grid = inp_dict["N_Grid"] R_sphere = inp_dict["R_sphere"] x_mod_arr_3d = COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict) WR_arr_3d = np.zeros((N_Grid,N_Grid,N_Grid)) bool_arr = x_mod_arr_3d <= R_sphere*2.0 WR_arr_3d[bool_arr] = np.exp(-1.0 * x_mod_arr_3d[bool_arr] / R_sphere) norm = np.sum(WR_arr_3d) WR_arr_3d /= norm return WR_arr_3d,0 def COMPUTE_SPHERICAL_TOP_HAT_WITH_SLICES_REAL_SPACE(inp_dict): """ Very Important Spherical top hat is calculated in terms of grid size and not comoving or physical distances. This makes easier to use same code for different box sizes. """ N_Grid = inp_dict["N_Grid"] R_sphere = inp_dict["R_sphere"] N_angle = inp_dict["N_angle"] polar_tup = COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX_POLAR_COORDINATES(inp_dict) R_arr_3d = polar_tup[0] theta_arr_3d = polar_tup[1] phi_arr_3d = polar_tup[2] WR_arr_3d = np.zeros((N_Grid,N_Grid,N_Grid)) bool_arr = R_arr_3d <= R_sphere npix = np.sum(bool_arr) V_sphere_inv = 1.0 / float(npix) # Normalization constant WR_arr_3d[bool_arr] = V_sphere_inv if N_angle > 1: sector_dict = GENERATE_SECTOR_INFO_DICTIONRAY(N_angle) N_sector = len(sector_dict) sector_info_arr = np.zeros(WR_arr_3d.shape,dtype = np.int32) for sector_indx in xrange(N_sector): theta_bound = sector_dict[sector_indx]["theta_bound"] phi_bound = sector_dict[sector_indx]["phi_bound"] bool_arr_theta = (theta_bound[0] <= theta_arr_3d) & (theta_arr_3d <= theta_bound[1]) # Note <= and < signs bool_arr_phi = (phi_bound[0] <= phi_arr_3d) & (phi_arr_3d <= phi_bound[1]) # Note <= and < signs bool_arr = bool_arr_theta & bool_arr_phi if np.sum(bool_arr) != 0: sector_info_arr[bool_arr] = sector_indx + 1 # Note +1 is intentional else: sector_info_arr = 0.0 return WR_arr_3d,npix,sector_info_arr def GENERATE_SECTOR_INFO_DICTIONRAY(N_angle): #theta_arr = np.linspace(0.0,np.pi,N_angle+1) #phi_arr = np.linspace(-np.pi,np.pi,N_angle+1) dtheta = np.pi * 0.05 dphi = 2.0*np.pi * 0.05 theta_arr = np.linspace(0.0-dtheta,np.pi+dtheta,N_angle+1) phi_arr = np.linspace(-np.pi-dphi,np.pi+dphi,N_angle+1) sector_dict = {} sector_indx = 0 for theta_indx in xrange(N_angle): theta_lb = theta_arr[theta_indx] theta_ub = theta_arr[theta_indx+1] theta_bound = [theta_lb,theta_ub] if (theta_indx > 0) & (theta_indx < N_angle -1): theta_bound = [theta_lb-dtheta,theta_ub+dtheta] for phi_indx in xrange(N_angle): phi_lb = phi_arr[phi_indx] phi_ub = phi_arr[phi_indx+1] phi_bound = [phi_lb,phi_ub] if (phi_indx > 0) & (phi_indx < N_angle -1): phi_bound = [phi_lb-dphi,phi_ub+dphi] sector_sub_dict = {sector_indx:{"theta_bound":theta_bound,"phi_bound":phi_bound}} sector_indx += 1 sector_dict.update(sector_sub_dict) return sector_dict def COMPUTE_SPHERICAL_TOP_HAT_WITH_SUBCUBES_REAL_SPACE(inp_dict): """ Very Important Spherical top hat is calculated in terms of grid size and not comoving or physical distances. This makes easier to use same code for different box sizes. """ N_Grid = inp_dict["N_Grid"] R_sphere = inp_dict["R_sphere"] N_subcube = 8 mesh_tup = COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict,"y") x_arr_3d = mesh_tup[0] y_arr_3d = mesh_tup[1] z_arr_3d = mesh_tup[2] R_arr_3d = mesh_tup[3] WR_arr_3d = np.zeros((N_Grid,N_Grid,N_Grid)) bool_arr = R_arr_3d <= R_sphere npix = np.sum(bool_arr) V_sphere_inv = 1.0 / float(npix) # Normalization constant WR_arr_3d[bool_arr] = V_sphere_inv if N_subcube > 1: sector_info_arr = GENERATE_SUBCUBE_INFO_DICTIONRAY(x_arr_3d,y_arr_3d,z_arr_3d,WR_arr_3d) else: sector_info_arr = 0.0 return WR_arr_3d,npix,sector_info_arr def GENERATE_SUBCUBE_INFO_DICTIONRAY(x_arr_3d,y_arr_3d,z_arr_3d,WR_arr_3d): sector_info_arr = np.zeros(WR_arr_3d.shape,dtype = np.int32) bool_arr_sphere = WR_arr_3d > 0.0 sector_indx = 1 for sign_val_x in [-1,1]: bool_arr_x = FIND_QUADRENT(x_arr_3d,sign_val_x) for sign_val_y in [-1,1]: bool_arr_y = FIND_QUADRENT(y_arr_3d,sign_val_y) for sign_val_z in [-1,1]: bool_arr_z = FIND_QUADRENT(z_arr_3d,sign_val_z) bool_arr = bool_arr_sphere & bool_arr_x & bool_arr_y & bool_arr_z sector_info_arr[bool_arr] = sector_indx sector_indx += 1 return sector_info_arr def FIND_QUADRENT(xyz_arr_3d,sign_val): if sign_val == -1: bool_arr_xyz = xyz_arr_3d <= 0 elif sign_val == 1: bool_arr_xyz = xyz_arr_3d >= 0 return bool_arr_xyz ####################################################################### ####### Compute Window function in Fourier Space MPI Routines ######### ####################################################################### def MPI_COMPUTE_WINDOW_FUNCTION_REAL_SPACE(inp_dict,flag_window="top_hat"): if flag_window == "top_hat": # Compute Spherical Top Hat Function WR_arr_3d,npix = MPI_COMPUTE_SPHERICAL_TOP_HAT_REAL_SPACE(inp_dict) return WR_arr_3d,npix def MPI_COMPUTE_SPHERICAL_TOP_HAT_REAL_SPACE(inp_dict): """ Very Important Spherical top hat is calculated in terms of grid size and not comoving or physical distances. This makes easier to use same code for different box sizes. """ from common_lib.MPI_functions.mpi_init import mpi_variable,\ mpi_check_number_of_processors mpi_dict = mpi_variable() comm = mpi_dict["comm"] rank = mpi_dict["rank"] N_proc = mpi_dict["N_proc"] N_Grid = inp_dict["N_Grid"] R_sphere = inp_dict["R_sphere"] N_Grid_per_proc = mpi_check_number_of_processors(N_Grid,N_proc) lb_indx = rank * N_Grid_per_proc ub_indx = (rank+1) * N_Grid_per_proc x_mod_arr_3d = MPI_COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict,lb_indx,ub_indx) WR_arr_3d = np.zeros((N_Grid_per_proc,N_Grid,N_Grid)) bool_arr = x_mod_arr_3d <= R_sphere npix_proc = np.sum(bool_arr) if rank == 0: npix = npix_proc for proc_indx in range(1,N_proc): npix_proc_indx = comm.recv(source=proc_indx,tag=proc_indx) npix = npix + npix_proc_indx comm.Barrier() for proc_indx in range(1,N_proc): comm.send(npix,dest=proc_indx,tag=proc_indx) comm.Barrier() else: comm.send(npix_proc,dest=0,tag=rank) comm.Barrier() npix = comm.recv(source=0,tag=rank) comm.Barrier() V_sphere_inv = 1.0 / float(npix) # Normalization constant WR_arr_3d[bool_arr] = V_sphere_inv return WR_arr_3d,npix def MPI_COMPUTE_3D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict,lb_indx,ub_indx): """ Compute 3D Cartesian Grid This routine computes 3D grid in integers like (-512,-511, ... ,0, ... ,510,511) """ if "x_arr_1d" not in inp_dict.keys(): inp_dict = COMPUTE_1D_REAL_SPACE_ARRAY_IN_BOX_GRID_INDEX(inp_dict) x_arr_1d = inp_dict["x_arr_1d"] x_arr_3d_tup = np.meshgrid(x_arr_1d[lb_indx:ub_indx],x_arr_1d,x_arr_1d,indexing = 'ij') # Note when 'ij' is used the order is correct x_arr_3d = x_arr_3d_tup[0] y_arr_3d = x_arr_3d_tup[1] z_arr_3d = x_arr_3d_tup[2] x_mod_arr_3d = (x_arr_3d**2.0 + y_arr_3d**2.0 + z_arr_3d**2.0)**0.5 return x_mod_arr_3d def CHECK_IF_KEY_PRESENT(key_list,inp_dict): inp_dict_key_list = list(inp_dict.keys()) flag_option = False for key_name_reqd in key_list: if key_name_reqd not in inp_dict_key_list: flag_option = True print("inp_dict should contain key :",key_name_reqd) if flag_option: print("Exiting the code now !!!") exit()