def advected_field(cat, field_late, weight, nmesh): ''' Get the advected component field with the weights given by the 'weight' column. Input: -cat: catalog of galaxy positions, assumed to have dimensions [3, N] right now -field_late: late-time density field (the delta_1 component) which is needed to prevent an nbodykit bug when placing down weights -weight: the early time component field used for the weights. See Note about 1-1 mapping -nmesh: Nmesh for the late-time advected field. Can be different from the grid value of the component field. Output: -field_complate: late-time, advected component field for the given input weights. Notes: In ZAemulus it's assumed that there's a clear 1-1 mapping between positions and particle since this comes from reshaping the grids. This has to be changed for a more realistic particle catalog. ''' Lbox = field_late.BoxSize[0] nbodycat = np.empty(nmesh**3, dtype=[('Position', ('f8', 3)), ('Weight', 'f8')]) nbodycat['Position'] = cat.T #Collapse the weight grids into catalog shape. #Need to change this for future runs. #+1 is needed if subtracting from field_late nbodycat['Weight'] = weight.value.reshape(nmesh**3) + 1 nbkcat = ArrayCatalog(nbodycat, Nmesh=nmesh, BoxSize=Lbox) mesh = nbkcat.to_mesh(Nmesh=nmesh, BoxSize=Lbox, weight='Weight') #Paint and subtract from late field to avoid annoying nbodykit bugs field_component_late = (mesh.paint(mode='real') - 1) field_component_late -= field_late return field_component_late
def __init__(self, ra, dec, collision_radius=62/60./60., seed=None, degrees=True, comm=None): # compute the pos ra = ArrayCatalog.make_column(ra) dec = ArrayCatalog.make_column(dec) pos = SkyToUnitSphere(ra, dec, degrees=degrees).compute() # make the source dt = numpy.dtype([('Position', (pos.dtype.str, 3))]) pos = numpy.squeeze(pos.view(dtype=dt)) source = ArrayCatalog(pos, BoxSize=numpy.array([2., 2., 2.]), comm=comm) self.source = source self.comm = source.comm # set the seed randomly if it is None if seed is None: if self.comm.rank == 0: seed = numpy.random.randint(0, 4294967295) seed = self.comm.bcast(seed) # save the attrs self.attrs = {} self.attrs['collision_radius'] = collision_radius self.attrs['seed'] = seed self.attrs['degrees'] = degrees # store collision radius in radians self._collision_radius_rad = numpy.deg2rad(collision_radius) if self.comm.rank == 0: self.logger.info("collision radius in degrees = %.4f" %collision_radius) # compute self.run()
def find_features(self, peakcolumn=None): """ Based on the particle labels, identify the groups, and return the center-of-mass ``CMPosition``, ``CMVelocity``, and Length of each feature. If a ``peakcolumn`` is given, the ``PeakPosition`` and ``PeakVelocity`` is also calculated for the particle at the peak value of the column. Data is scattered evenly across all ranks. Returns ------- :class:`~nbodykit.source.catalog.array.ArrayCatalog` : a source holding the ('CMPosition', 'CMVelocity', 'Length') of each feature; optionaly, ``PeakPosition``, ``PeakVelocity`` are also included if ``peakcolumn`` is not None """ # the center-of-mass (Position, Velocity, Length) halos = fof_catalog(self._source, self.labels, self.comm, peakcolumn=peakcolumn, periodic=self.attrs['periodic']) attrs = self._source.attrs.copy() attrs.update(self.attrs) return ArrayCatalog(halos, comm=self.comm, **attrs)
def cat_to_mesh(cat, Lbox, Nmesh): ''' Convert the late-time particle catalog to a friendly neighbourhood mesh Input: -cat: format is [Npos, Nparticles] -Lbox: units are Mpc/h -Nmesh: size of mesh you want to deposit particles in Output: -mesh ''' nbodycat = np.empty(nmesh**3, dtype=[('Position', ('f8', 3))]) nbodycat['Position'] = cat.T nbkcat = ArrayCatalog(nbodycat, Nmesh=nmesh, BoxSize=Lbox) mesh = nbkcat.to_mesh(Nmesh=nmesh, BoxSize=Lbox) field = (mesh.paint(mode='real') - 1) return field
def __init__(self, ra, dec, collision_radius=62 / 60. / 60., seed=None, degrees=True, comm=None): # compute the pos ra = ArrayCatalog.make_column(ra) dec = ArrayCatalog.make_column(dec) pos = SkyToUnitSphere(ra, dec, degrees=degrees).compute() # make the source dt = numpy.dtype([('Position', (pos.dtype.str, 3))]) pos = numpy.squeeze(pos.view(dtype=dt)) source = ArrayCatalog(pos, BoxSize=numpy.array([2., 2., 2.]), comm=comm) self.source = source self.comm = source.comm # set the seed randomly if it is None if seed is None: if self.comm.rank == 0: seed = numpy.random.randint(0, 4294967295) seed = self.comm.bcast(seed) # save the attrs self.attrs = {} self.attrs['collision_radius'] = collision_radius self.attrs['seed'] = seed self.attrs['degrees'] = degrees # store collision radius in radians self._collision_radius_rad = numpy.deg2rad(collision_radius) if self.comm.rank == 0: self.logger.info("collision radius in degrees = %.4f" % collision_radius) # compute self.run()
def ArrayCatalogExample(): from nbodykit.source.catalog import ArrayCatalog print("https://nbodykit.readthedocs.io/en/latest/catalogs/reading.html#array-data") # generate the fake data num_data = 4096*4096 BoxSize = 1380 data = numpy.empty(num_data, dtype=[('Position', ('f8', 3))]) data['Position'] = numpy.random.uniform(size=(num_data, 3)) * BoxSize # save to a npy file numpy.save("npy-example.npy", data) data = numpy.load("npy-example.npy") # initialize the catalog cat = ArrayCatalog(data, BoxSize=BoxSize, Nmesh=128) print(cat) AnalyzeCatalog(cat)
def main(): """ Script to convert Rockstar halo catalog to bigfile catalog. Login to a single node on helios and run there on command line. For batch runs, use e.g. for SEED in {0..1}; do python main_rockstar_catalog_to_bigfile.py --rockstar_halos "/scratch/mschmittfull/lss/ms_gadget/run4/0000040${SEED}-01536-1500.0-wig/snap_0.6250.gadget3/rockstar_out_0.list" --max_rows 5 --include_parent_ID; done """ setup_logging() ap = ArgumentParser() ap.add_argument( '--rockstar_halos', help= ('File name of Rockstar halo catalog, e.g.' '/data/mschmittfull/lss/ms_gadget/run4/00000400-01536-500.0-wig/snap_0.6250.gadget3/rockstar_out_0.list' ), default= '/scratch/mschmittfull/lss/ms_gadget/run4/00000400-01536-500.0-wig/snap_0.6250.gadget3/rockstar_out_0.list.parents' ) ap.add_argument('--add_RSD', dest='RSD', action='store_true', help='Add RSD to position') ap.add_argument('--include_parent_ID', dest='include_parent_ID', action='store_true', help='Include ID and parent ID in bigfile.') # ap.add_argument( # '--RSD', help='Add RSD to positions if not 0', # type=int, # default=0) ap.add_argument('--max_rows', help='Max number of rows to read. Read all if 0.', type=int, default=0) ap.set_defaults(RSD=False, include_parent_ID=False) args = ap.parse_args() RSD_LOS = np.array([0, 0, 1]) # load input halo catalog print('Read halos from %s' % args.rockstar_halos) # read header with open(args.rockstar_halos) as myfile: header = [next(myfile) for x in xrange(16)] header = ''.join(header) print('Header:') print(header) # get names of columns np_cat1 = np.genfromtxt(args.rockstar_halos, names=True, max_rows=1) names = np_cat1.dtype.names # keep only a subset usecol_names = ['X', 'Y', 'Z', 'VX', 'VY', 'VZ', 'Mvir'] if args.include_parent_ID: usecol_names += ['ID', 'PID'] usecols = [] for column_number, name in enumerate(names): if name in usecol_names: usecols.append(column_number) print('usecols:', usecols) print([names[usecol] for usecol in usecols]) # read data. print('Reading data') if args.max_rows == 0: max_rows = None else: max_rows = args.max_rows # TODO: np.loadtxt should be faster, but now take 5 minutes so probably ok. np_cat = np.genfromtxt(args.rockstar_halos, names=True, max_rows=max_rows, usecols=usecols) print('Read data:') print(np_cat[:5]) # convert to arraycatalog cat = ArrayCatalog(np_cat) # fill position and velocity pos = np.empty(cat.csize, dtype=[('Position', ('f8', 3))]) pos['Position'][:, 0] = cat['X'] pos['Position'][:, 1] = cat['Y'] pos['Position'][:, 2] = cat['Z'] cat['Position'] = pos['Position'] del pos vel = np.empty(cat.csize, dtype=[('Velocity', ('f8', 3))]) vel['Velocity'][:, 0] = cat['VX'] vel['Velocity'][:, 1] = cat['VY'] vel['Velocity'][:, 2] = cat['VZ'] # todo: what units? cat['Velocity'] = vel['Velocity'] del vel cat['log10Mvir'] = np.log10(cat['Mvir']) # Keep only some columns keep_columns = ['Position', 'Velocity', 'log10Mvir'] if args.include_parent_ID: # also keep halo ID and parent ID keep_columns += ['ID', 'PID'] cat = catalog_persist(cat, keep_columns) cat.attrs['rockstar_header'] = header if args.RSD: raise Exception('RSD not implemented') print('Will write data:') for c in keep_columns: print('%s:' % c, cat[c]) # save to bigfile if max_rows is None: out_fname = '%s.bigfile' % args.rockstar_halos else: out_fname = '%s_max_rows%d.bigfile' % (args.rockstar_halos, max_rows) if os.path.exists(out_fname): rmtree(out_fname) if cat.comm.rank == 0: print('Writing to %s' % out_fname) cat.save(out_fname, columns=keep_columns) if cat.comm.rank == 0: print('Wrote %s' % out_fname)
def run(self): """ Compute the cylindrical groups, saving the results to the :attr:`groups` attribute Attributes ---------- groups : :class:`~nbodykit.source.catalog.array.ArrayCatalog` a catalog holding the result of the grouping. The length of the catalog is equal to the length of the input size, i.e., the length is equal to the :attr:`size` attribute. The relevant fields are: #. cgm_type : a flag specifying the type for each object, with 0 specifying CGM central and 1 denoting CGM satellite #. cgm_haloid : The index of the CGM object this object belongs to; an integer between 0 and the total number of CGM halos #. num_cgm_sats : The number of satellites in the CGM halo """ from pmesh.domain import GridND from nbodykit.algorithms.fof import split_size_3d comm = self.comm rperp, rpar = self.attrs['rperp'], self.attrs['rpar'] rankby = self.attrs['rankby'] if self.attrs['periodic']: boxsize = self.attrs['BoxSize'] else: boxsize = None np = split_size_3d(self.comm.size) if self.comm.rank == 0: self.logger.info("using cpu grid decomposition: %s" % str(np)) # add a column for original index self.source['origind'] = self.source.Index # sort the data data = self.source.sort(self.attrs['rankby'], usecols=['Position', 'origind']) # add a column to track sorted index data['sortindex'] = data.Index # global min/max across all ranks pos = data.compute(data['Position']) posmin = numpy.asarray(comm.allgather(pos.min(axis=0))).min(axis=0) posmax = numpy.asarray(comm.allgather(pos.max(axis=0))).max(axis=0) # domain decomposition grid = [ numpy.linspace(posmin[0], posmax[0], np[0] + 1, endpoint=True), numpy.linspace(posmin[0], posmax[1], np[1] + 1, endpoint=True), numpy.linspace(posmin[0], posmax[2], np[2] + 1, endpoint=True), ] domain = GridND(grid, comm=comm) # run the CGM algorithm groups = cgm(comm, data, domain, rperp, rpar, self.attrs['flat_sky_los'], boxsize) # make the final structured array self.groups = ArrayCatalog(groups, comm=self.comm, **self.attrs) # log some info N_cen = (groups['cgm_type'] == 0).sum() isolated_N_cen = ((groups['cgm_type'] == 0) & (groups['num_cgm_sats'] == 0)).sum() N_cen = self.comm.allreduce(N_cen) isolated_N_cen = self.comm.allreduce(isolated_N_cen) if self.comm.rank == 0: self.logger.info("found %d CGM centrals total" % N_cen) self.logger.info("%d/%d are isolated centrals (no satellites)" % (isolated_N_cen, N_cen)) # delete the column we added to source del self.source['origind']
def to_halos(self, particle_mass, cosmo, redshift, mdef='vir', posdef='cm', peakcolumn='Density'): """ Return a :class:`~nbodykit.source.catalog.halos.HaloCatalog`, holding the center-of-mass position and velocity of each FOF halo, as well as the properly scaled mass, for a given cosmology and redshift. The returned catalog also has default analytic prescriptions for halo radius and concentration. The data is scattered evenly across all ranks. Parameters ---------- particle_mass : float the particle mass, which is used to convert the number of particles in each halo to a total mass cosmo : :class:`nbodykit.cosmology.core.Cosmology` the cosmology of the catalog redshift : float the redshift of the catalog mdef : str, optional string specifying mass definition, used for computing default halo radii and concentration; should be 'vir' or 'XXXc' or 'XXXm' where 'XXX' is an int specifying the overdensity posdef : str, optional position, can be cm (center of mass) or peak (particle with maximum value on a column) peakcolumn : str , optional when posdef is 'peak', this is the column in source for identifying particles at the peak for the position and velocity. Returns ------- :class:`~nbodykit.source.catalog.halos.HaloCatalog` a HaloCatalog at the specified cosmology and redshift """ from nbodykit.source.catalog.halos import HaloCatalog assert posdef in ['cm', 'peak'], "``posdef`` should be 'cm' or 'peak'" # meta-data attrs = self._source.attrs.copy() attrs.update(self.attrs) attrs['particle_mass'] = particle_mass if posdef == 'cm': # using the center-of-mass (Position, Velocity, Length) for each halo # not needing a column for peaks. peakcolumn = None else: pass data = fof_catalog(self._source, self.labels, self.comm, peakcolumn=peakcolumn, periodic=self.attrs['periodic']) data = data[data['Length'] > 0] halos = ArrayCatalog(data, comm=self.comm, **attrs) if posdef == 'cm': halos['Position'] = halos['CMPosition'] halos['Velocity'] = halos['CMVelocity'] elif posdef == 'peak': halos['Position'] = halos['PeakPosition'] halos['Velocity'] = halos['PeakVelocity'] # add the halo mass column halos['Mass'] = particle_mass * halos['Length'] coldefs = { 'mass': 'Mass', 'velocity': 'Velocity', 'position': 'Position' } return HaloCatalog(halos, cosmo, redshift, mdef=mdef, **coldefs)
y = np.loadtxt('../../../nb/data/y.txt') z = np.loadtxt('../../../nb/data/z.txt') mass = np.loadtxt('../../../nb/data/m_c200.txt') length = len(x) data = np.empty(length, dtype=[('Position', ('f8', 3)), ('Mass', 'f8')]) pos = np.zeros((length, 3)) for i in np.arange(length): pos[i][0] = x[i] pos[i][1] = y[i] pos[i][2] = z[i] data['Position'] = pos data['Mass'] = mass # initialize the catalog f = ArrayCatalog(data) print(f) print("columns = ", f.columns) # default Weight,Selection also present print("total size = ", f.csize) f = ArrayCatalog({'Position': data['Position'], 'Mass': data['Mass']}) print(f) print("columns = ", f.columns) # default Weight,Selection also present print("total size = ", f.csize) # convert to a MeshSource, using CIC interpolation on 1280^3 mesh mesh = f.to_mesh(window='cic', Nmesh=1280, BoxSize=4000.0,
def main(): """ Convert grids4plots 3d grids to 2d slices. """ ##################################### # PARSE COMMAND LINE ARGS ##################################### ap = ArgumentParser() ap.add_argument('--inbasepath', type=str, default='$SCRATCH/perr/grids4plots/', help='Input base path.') ap.add_argument('--outbasepath', type=str, default='$SCRATCH/perr/grids4plots/', help='Output base path.') ap.add_argument( '--inpath', type=str, default= 'main_calc_Perr_2020_Sep_22_18:44:31_time1600800271.dill', # laptop #default='main_calc_Perr_2020_Aug_26_02:49:57_time1598410197.dill', # cluster help='Input path.') ap.add_argument('--Rsmooth', type=float, default=2.0, help='3D Gaussian smoothing applied to field.') # min and max index included in output. inclusive. ap.add_argument('--ixmin', type=int, default=5, help='xmin of output. must be between 0 and Ngrid.') ap.add_argument('--ixmax', type=int, default=5, help='xmax of output') ap.add_argument('--iymin', type=int, default=0, help='ymin of output') ap.add_argument('--iymax', type=int, default=-1, help='ymax of output') ap.add_argument('--izmin', type=int, default=0, help='zmin of output') ap.add_argument('--izmax', type=int, default=-1, help='zmax of output') cmd_args = ap.parse_args() verbose = True ##################################### # START PROGRAM ##################################### comm = CurrentMPIComm.get() rank = comm.rank path = os.path.join(os.path.expandvars(cmd_args.inbasepath), cmd_args.inpath) if rank == 0: print('path: ', path) initialized_slicecat = False for fname in os.listdir(path): if fname.startswith('SLICE'): continue full_fname = os.path.join(path, fname) print('%d Reading %s' % (rank, full_fname)) inmesh = BigFileMesh(full_fname, dataset='tmp4storage', header='header') Ngrid = inmesh.attrs['Ngrid'] boxsize = inmesh.attrs['boxsize'] # apply smoothing mesh = apply_smoothing(inmesh, mode='Gaussian', R=cmd_args.Rsmooth) del inmesh # convert indices to modulo ngrid ixmin = cmd_args.ixmin % Ngrid ixmax = cmd_args.ixmax % Ngrid iymin = cmd_args.iymin % Ngrid iymax = cmd_args.iymax % Ngrid izmin = cmd_args.izmin % Ngrid izmax = cmd_args.izmax % Ngrid # convert to boxsize units (Mpc/h) xmin = float(ixmin) / float(Ngrid) * boxsize xmax = float(ixmax + 1) / float(Ngrid) * boxsize ymin = float(iymin) / float(Ngrid) * boxsize ymax = float(iymax + 1) / float(Ngrid) * boxsize zmin = float(izmin) / float(Ngrid) * boxsize zmax = float(izmax + 1) / float(Ngrid) * boxsize if not initialized_slicecat: # Generate catalog with positions of slice points, to readout mesh there. # First generate all 3D points. Then keep only subset in slice. # THen readout mesh at those points. # use pmesh generate_uniform_particle_grid # http://rainwoodman.github.io/pmesh/pmesh.pm.html?highlight= # readout#pmesh.pm.ParticleMesh.generate_uniform_particle_grid partmesh = ParticleMesh(BoxSize=boxsize, Nmesh=[Ngrid, Ngrid, Ngrid]) ptcles = partmesh.generate_uniform_particle_grid(shift=0.0, dtype='f8') #print("type ptcles", type(ptcles), ptcles.shape) #print("head ptcles:", ptcles[:5,:]) dtype = np.dtype([('Position', ('f8', 3))]) # number of rows is given by number of ptcles on this rank uni_cat_array = np.empty((ptcles.shape[0], ), dtype=dtype) uni_cat_array['Position'] = ptcles uni_cat = ArrayCatalog(uni_cat_array, comm=None, BoxSize=boxsize * np.ones(3), Nmesh=[Ngrid, Ngrid, Ngrid]) del ptcles del uni_cat_array print("%d: Before cut: local Nptcles=%d, global Nptcles=%d" % (comm.rank, uni_cat.size, uni_cat.csize)) # only keep points in the slice uni_cat = uni_cat[(uni_cat['Position'][:, 0] >= xmin) & (uni_cat['Position'][:, 0] < xmax) & (uni_cat['Position'][:, 1] >= ymin) & (uni_cat['Position'][:, 1] < ymax) & (uni_cat['Position'][:, 2] >= zmin) & (uni_cat['Position'][:, 2] < zmax)] print("%d: After cut: local Nptcles=%d, global Nptcles=%d" % (comm.rank, uni_cat.size, uni_cat.csize)) initialized_slicecat = True # read out full 3D mesh at catalog positions. this is a numpy array slicecat = readout_mesh_at_cat_pos(mesh=mesh, cat=uni_cat, readout_window='nearest') if rank == 0: print('slicecat type:', type(slicecat)) slicecat = GatherArray(slicecat, comm, root=0) if rank == 0: if not slicecat.shape == ((ixmax - ixmin + 1) * (iymax - iymin + 1) * (izmax - izmin + 1), ): raise Exception( 'Unexpected shape of particles read out on slice: %s' % str(slicecat.shape)) slicecat = slicecat.reshape( (ixmax - ixmin + 1, iymax - iymin + 1, izmax - izmin + 1)) print('slicecat shape:', slicecat.shape) if verbose: print('slicecat:', slicecat) # convert to a mesh. assume full numpy array sits on rank 0. Lx = xmax - xmin Ly = ymax - ymin Lz = zmax - zmin if Lx == 0.: Lx = boxsize / float(Ngrid) if Ly == 0.: Ly = boxsize / float(Ngrid) if Lz == 0.: Lz = boxsize / float(Ngrid) BoxSize_slice = np.array([Lx, Ly, Lz]) slicemesh = ArrayMesh(slicecat, BoxSize=BoxSize_slice, root=0) outshape = slicemesh.compute(mode='real').shape if verbose: print('slicemesh: ', slicemesh.compute(mode='real')) # write to disk outpath = os.path.join( os.path.expandvars(cmd_args.outbasepath), cmd_args.inpath, 'SLICE_R%g_%d-%d_%d-%d_%d-%d/' % (cmd_args.Rsmooth, ixmin, ixmax, iymin, iymax, izmin, izmax)) if rank == 0: if not os.path.exists(outpath): os.makedirs(outpath) full_outfname = os.path.join(outpath, fname) if rank == 0: print('Writing %s' % full_outfname) slicemesh.save(full_outfname) if rank == 0: print('Wrote %s' % full_outfname)
def run(self): """ Run the fiber assignment algorithm. This attaches the following attribute to the object: - :attr:`labels` .. note:: The :attr:`labels` attribute has a 1-to-1 correspondence with the rows in the input source. Attributes ---------- labels: :class:`~nbodykit.source.catalog.array.ArrayCatalog`; size: :attr:`size` a CatalogSource that has the following columns: - Label : the group labels for each object in the input CatalogSource; label == 0 objects are not in a group - Collided : a flag array specifying which objects are collided, i.e., do not receive a fiber - NeighborID : for those objects that are collided, this gives the (global) index of the nearest neighbor on the sky (0-indexed) in the input catalog ``source``, else it is set to -1 """ from astropy.utils.misc import NumpyRNGContext from nbodykit.algorithms import FOF # angular FOF: labels gives the global group ID corresponding to each # object in Position on this rank fof = FOF(self.source, self._collision_radius_rad, 1, absolute=True) # assign the fibers (in parallel) with NumpyRNGContext(self.attrs['seed']): collided, neighbors = self._assign_fibers(fof.labels) # all reduce to get summary statistics N_pop1 = self.comm.allreduce((collided ^ 1).sum()) N_pop2 = self.comm.allreduce((collided).sum()) f = N_pop2 * 1. / (N_pop1 + N_pop2) # print out some info if self.comm.rank == 0: self.logger.info("population 1 (clean) size = %d" % N_pop1) self.logger.info("population 2 (collided) size = %d" % N_pop2) self.logger.info("collision fraction = %.4f" % f) # return a structured array d = list( zip(['Label', 'Collided', 'NeighborID'], [fof.labels, collided, neighbors])) dtype = numpy.dtype([(col, x.dtype) for col, x in d]) result = numpy.empty(len(fof.labels), dtype=dtype) for col, x in d: result[col] = x # make a particle source self.labels = ArrayCatalog(result, comm=self.comm, **self.source.attrs)