def to_real_field(self, out=None, normalize=True): """ Paint the density field, by interpolating the position column on to the mesh. This computes the following meta-data attributes in the process of painting, returned in the :attr:`attrs` attributes of the returned RealField object: - N : int the (unweighted) total number of objects painted to the mesh - W : float the weighted number of total objects, equal to the collective sum of the 'weight' column - shotnoise : float the Poisson shot noise, equal to the volume divided by ``N`` - num_per_cell : float the mean number of weighted objects per cell .. note:: The density field on the mesh is normalized as :math:`1+\delta`, such that the collective mean of the field is unity. See the :ref:`documentation <painting-mesh>` on painting for more details on painting catalogs to a mesh. Returns ------- real : :class:`pmesh.pm.RealField` the painted real field; this has a ``attrs`` dict storing meta-data """ pm = self.pm Nlocal = 0 # (unweighted) number of particles read on local rank Wlocal = 0 # (weighted) number of particles read on local rank W2local = 0 # sum of weight square. This is used to estimate shotnoise. # the paint brush window resampler = window.methods[self.resampler] # initialize the RealField to return if out is not None: assert isinstance(out, RealField), "output of to_real_field must be a RealField" numpy.testing.assert_array_equal(out.pm.Nmesh, pm.Nmesh) toret = out else: toret = RealField(pm) toret[:] = 0 # for interlacing, we need two empty meshes if out was provided # since out may have non-zero elements, messing up our interlacing sum if self.interlaced: real1 = RealField(pm) real1[:] = 0 # the second, shifted mesh (always needed) real2 = RealField(pm) real2[:] = 0 Position = self.Position Weight = self.Weight Value = self.Value Selection = self.Selection # ensure the slices are synced, since decomposition is collective Nlocalmax = max(pm.comm.allgather(len(Position))) H = pm.BoxSize / pm.Nmesh # paint data in chunks on each rank; # we do this by chunk 8 million is pretty big anyways. max_chunksize = _global_options['paint_chunk_size'] # use a local scope to avoid having two copies of data in memory def dochunk(s): if len(Position) != 0: # selection has to be computed many times when data is `large`. columns = [Position[s]] if Weight is not None: columns.append(Weight[s]) if Value is not None: columns.append(Value[s]) if Selection is not None: columns.append(Selection[s]) # be sure to use the source to compute data = self.source.compute(columns) sel = Ellipsis if Selection is None else data.pop() value = None if Value is None else data.pop()[sel] weight = None if Weight is None else data.pop()[sel] position = data.pop()[sel] else: # workaround a potential dask issue on empty dask arrays position = numpy.empty((0, 3), dtype=Position.dtype) weight = None value = None if weight is None: weight = numpy.ones(len(position)) if value is None: value = numpy.ones(len(position)) # track total (selected) number and sum of weights Nlocal = len(position) Wlocal = weight.sum() W2local = (weight ** 2).sum() # no interlacing if not self.interlaced: lay = pm.decompose(position, smoothing=0.5 * resampler.support) else: lay = pm.decompose(position, smoothing=1.0 * resampler.support) # if we are receiving too many particles, abort and retry with a smaller chunksize newlengths = pm.comm.allgather(lay.newlength) if any([newlength > 2 * max_chunksize for newlength in newlengths]): if pm.comm.rank == 0: self.logger.info("Throttling chunksize as some ranks will receive too many particles. (%d > %d)" % (max(newlengths), max_chunksize * 2)) raise StopIteration p = lay.exchange(position) w = lay.exchange(weight) v = lay.exchange(value) if not self.interlaced: pm.paint(p, mass=w * v, resampler=resampler, hold=True, out=toret) # interlacing: use 2 meshes separated by 1/2 cell size else: # in mesh units shifted = pm.affine.shift(0.5) # paint to two shifted meshes pm.paint(p, mass=w * v, resampler=resampler, hold=True, out=real1) pm.paint(p, mass=w * v, resampler=resampler, transform=shifted, hold=True, out=real2) return Nlocal, Wlocal, W2local import gc i = 0 chunksize = max_chunksize while i < Nlocalmax: s = slice(i, i + chunksize) if pm.comm.rank == 0: self.logger.info("Chunk %d ~ %d / %d " % (i, i + chunksize, Nlocalmax)) try: Nlocal1, Wlocal1, W2local1 = dochunk(s) chunksize = min(max_chunksize, int(chunksize * 1.5)) except StopIteration: chunksize = chunksize // 2 if chunksize < 1: raise RuntimeError("Cannot find a chunksize that fits into memory.") continue finally: # collect unfreed items gc.collect() Nlocal += Nlocal1 Wlocal += Wlocal1 W2local += W2local1 Nglobal = pm.comm.allreduce(Nlocal) if pm.comm.rank == 0: self.logger.info("painted %d out of %d objects to mesh" % (Nglobal, self.source.csize)) i = i + chunksize # now the loop over particles is done if not self.interlaced: # nothing to do, toret is already filled. pass else: # compose the two interlaced fields into the final result. c1 = real1.r2c() c2 = real2.r2c() # and then combine for k, s1, s2 in zip(c1.slabs.x, c1.slabs, c2.slabs): kH = sum(k[i] * H[i] for i in range(3)) s1[...] = s1[...] * 0.5 + s2[...] * 0.5 * numpy.exp(0.5 * 1j * kH) # FFT back to real-space # NOTE: cannot use "toret" here in case user supplied "out" c1.c2r(real1) # need to add to the returned mesh if user supplied "out" toret[:] += real1[:] # unweighted number of objects N = pm.comm.allreduce(Nlocal) # weighted number of objects W = pm.comm.allreduce(Wlocal) # weighted number of objects W2 = pm.comm.allreduce(W2local) # weighted number density (objs/cell) nbar = 1. * W / numpy.prod(pm.Nmesh) # make sure we painted something or nbar is nan; in which case # we set the density to uniform everywhere. if N == 0: warnings.warn(("trying to paint particle source to mesh, " "but no particles were found!"), RuntimeWarning ) # shot noise is volume / un-weighted number shotnoise = numpy.prod(pm.BoxSize) * W2 / W ** 2 # save some meta-data toret.attrs = {} toret.attrs['shotnoise'] = shotnoise toret.attrs['N'] = N toret.attrs['W'] = W toret.attrs['W2'] = W toret.attrs['num_per_cell'] = nbar csum = toret.csum() if pm.comm.rank == 0: self.logger.info("painted %d out of %d objects to mesh" %(N, self.source.csize)) self.logger.info("mean particles per cell is %g", nbar) self.logger.info("sum is %g ", csum) if normalize: if nbar > 0: toret[...] /= nbar else: toret[...] = 1 if pm.comm.rank == 0: self.logger.info("normalized the convention to 1 + delta") return toret
def to_real_field(self, out=None, normalize=True): """ Paint the density field, by interpolating the position column on to the mesh. This computes the following meta-data attributes in the process of painting, returned in the :attr:`attrs` attributes of the returned RealField object: - N : int the (unweighted) total number of objects painted to the mesh - W : float the weighted number of total objects, equal to the collective sum of the 'weight' column - shotnoise : float the Poisson shot noise, equal to the volume divided by ``N`` - num_per_cell : float the mean number of weighted objects per cell .. note:: The density field on the mesh is normalized as :math:`1+\delta`, such that the collective mean of the field is unity. See the :ref:`documentation <painting-mesh>` on painting for more details on painting catalogs to a mesh. Returns ------- real : :class:`pmesh.pm.RealField` the painted real field; this has a ``attrs`` dict storing meta-data """ pm = self.pm Nlocal = 0 # (unweighted) number of particles read on local rank Wlocal = 0 # (weighted) number of particles read on local rank W2local = 0 # sum of weight square. This is used to estimate shotnoise. # the paint brush window resampler = window.methods[self.resampler] # initialize the RealField to return if out is not None: assert isinstance( out, RealField), "output of to_real_field must be a RealField" numpy.testing.assert_array_equal(out.pm.Nmesh, pm.Nmesh) toret = out else: toret = RealField(pm) toret[:] = 0 # for interlacing, we need two empty meshes if out was provided # since out may have non-zero elements, messing up our interlacing sum if self.interlaced: real1 = RealField(pm) real1[:] = 0 # the second, shifted mesh (always needed) real2 = RealField(pm) real2[:] = 0 Position = self.Position Weight = self.Weight Value = self.Value Selection = self.Selection # ensure the slices are synced, since decomposition is collective Nlocalmax = max(pm.comm.allgather(len(Position))) H = pm.BoxSize / pm.Nmesh # paint data in chunks on each rank; # we do this by chunk 8 million is pretty big anyways. max_chunksize = _global_options['paint_chunk_size'] # use a local scope to avoid having two copies of data in memory def dochunk(s): if len(Position) != 0: # selection has to be computed many times when data is `large`. columns = [Position[s]] if Weight is not None: columns.append(Weight[s]) if Value is not None: columns.append(Value[s]) if Selection is not None: columns.append(Selection[s]) # be sure to use the source to compute data = self.source.compute(columns) sel = Ellipsis if Selection is None else data.pop() value = None if Value is None else data.pop()[sel] weight = None if Weight is None else data.pop()[sel] position = data.pop()[sel] else: # workaround a potential dask issue on empty dask arrays position = numpy.empty((0, 3), dtype=Position.dtype) weight = None value = None if weight is None: weight = numpy.ones(len(position)) if value is None: value = numpy.ones(len(position)) # track total (selected) number and sum of weights Nlocal = len(position) Wlocal = weight.sum() W2local = (weight**2).sum() # no interlacing if not self.interlaced: lay = pm.decompose(position, smoothing=0.5 * resampler.support) else: lay = pm.decompose(position, smoothing=1.0 * resampler.support) # if we are receiving too many particles, abort and retry with a smaller chunksize recvlengths = pm.comm.allgather(lay.recvlength) if any( [recvlength > 2 * max_chunksize for recvlength in recvlengths]): if pm.comm.rank == 0: self.logger.info( "Throttling chunksize as some ranks will receive too many particles. (%d > %d)" % (max(recvlengths), max_chunksize * 2)) raise StopIteration p = lay.exchange(position) w = lay.exchange(weight) v = lay.exchange(value) if not self.interlaced: pm.paint(p, mass=w * v, resampler=resampler, hold=True, out=toret) # interlacing: use 2 meshes separated by 1/2 cell size else: # in mesh units shifted = pm.affine.shift(0.5) # paint to two shifted meshes pm.paint(p, mass=w * v, resampler=resampler, hold=True, out=real1) pm.paint(p, mass=w * v, resampler=resampler, transform=shifted, hold=True, out=real2) return Nlocal, Wlocal, W2local import gc i = 0 chunksize = max_chunksize while i < Nlocalmax: s = slice(i, i + chunksize) if pm.comm.rank == 0: self.logger.info("Chunk %d ~ %d / %d " % (i, i + chunksize, Nlocalmax)) try: Nlocal1, Wlocal1, W2local1 = dochunk(s) except StopIteration: chunksize = chunksize // 2 if chunksize < 1: raise RuntimeError( "Cannot find a chunksize that fits into memory.") continue finally: # collect unfreed items gc.collect() Nlocal += Nlocal1 Wlocal += Wlocal1 W2local += W2local1 Nglobal = pm.comm.allreduce(Nlocal) if pm.comm.rank == 0: self.logger.info("painted %d out of %d objects to mesh" % (Nglobal, self.source.csize)) i = i + chunksize chunksize = min(max_chunksize, int(chunksize * 1.5)) # now the loop over particles is done if not self.interlaced: # nothing to do, toret is already filled. pass else: # compose the two interlaced fields into the final result. c1 = real1.r2c() c2 = real2.r2c() # and then combine for k, s1, s2 in zip(c1.slabs.x, c1.slabs, c2.slabs): kH = sum(k[i] * H[i] for i in range(3)) s1[...] = s1[...] * 0.5 + s2[...] * 0.5 * numpy.exp( 0.5 * 1j * kH) # FFT back to real-space # NOTE: cannot use "toret" here in case user supplied "out" c1.c2r(real1) # need to add to the returned mesh if user supplied "out" toret[:] += real1[:] # unweighted number of objects N = pm.comm.allreduce(Nlocal) # weighted number of objects W = pm.comm.allreduce(Wlocal) # weighted number of objects W2 = pm.comm.allreduce(W2local) # weighted number density (objs/cell) nbar = 1. * W / numpy.prod(pm.Nmesh) # make sure we painted something or nbar is nan; in which case # we set the density to uniform everywhere. if N == 0: warnings.warn(("trying to paint particle source to mesh, " "but no particles were found!"), RuntimeWarning) # shot noise is volume / un-weighted number shotnoise = numpy.prod(pm.BoxSize) * W2 / W**2 # save some meta-data toret.attrs = {} toret.attrs['shotnoise'] = shotnoise toret.attrs['N'] = N toret.attrs['W'] = W toret.attrs['W2'] = W toret.attrs['num_per_cell'] = nbar csum = toret.csum() if pm.comm.rank == 0: self.logger.info("painted %d out of %d objects to mesh" % (N, self.source.csize)) self.logger.info("mean particles per cell is %g", nbar) self.logger.info("sum is %g ", csum) if normalize: if nbar > 0: toret[...] /= nbar else: toret[...] = 1 if pm.comm.rank == 0: self.logger.info("normalized the convention to 1 + delta") return toret
def to_real_field(self, out=None, normalize=True): """ Paint the density field, by interpolating the position column on to the mesh. This computes the following meta-data attributes in the process of painting, returned in the :attr:`attrs` attributes of the returned RealField object: - N : int the (unweighted) total number of objects painted to the mesh - W : float the weighted number of total objects, equal to the collective sum of the 'weight' column - shotnoise : float the Poisson shot noise, equal to the volume divided by ``N`` - num_per_cell : float the mean number of weighted objects per cell .. note:: The density field on the mesh is normalized as :math:`1+\delta`, such that the collective mean of the field is unity. See the :ref:`documentation <painting-mesh>` on painting for more details on painting catalogs to a mesh. Returns ------- real : :class:`pmesh.pm.RealField` the painted real field; this has a ``attrs`` dict storing meta-data """ # check for 'Position' column if self.position not in self.source: msg = "in order to paint a CatalogSource to a RealField, add a " msg += "column named '%s', representing the particle positions" % self.position raise ValueError(msg) pm = self.pm Nlocal = 0 # (unweighted) number of particles read on local rank Wlocal = 0 # (weighted) number of particles read on local rank # the paint brush window paintbrush = window.methods[self.window] # initialize the RealField to return if out is not None: assert isinstance( out, RealField), "output of to_real_field must be a RealField" numpy.testing.assert_array_equal(out.pm.Nmesh, pm.Nmesh) toret = out else: toret = RealField(pm) toret[:] = 0 # for interlacing, we need two empty meshes if out was provided # since out may have non-zero elements, messing up our interlacing sum if self.interlaced: real1 = RealField(pm) real1[:] = 0 # the second, shifted mesh (always needed) real2 = RealField(pm) real2[:] = 0 # read the necessary data (as dask arrays) columns = [self.position, self.weight, self.value, self.selection] Position, Weight, Value, Selection = self.source.read(columns) # ensure the slices are synced, since decomposition is collective Nlocalmax = max(pm.comm.allgather(len(Position))) # paint data in chunks on each rank; # we do this by chunk 8 million is pretty big anyways. chunksize = _global_options['paint_chunk_size'] for i in range(0, Nlocalmax, chunksize): s = slice(i, i + chunksize) if len(Position) != 0: # selection has to be computed many times when data is `large`. sel = self.source.compute(Selection[s]) # be sure to use the source to compute position, weight, value = \ self.source.compute(Position[s], Weight[s], Value[s]) # FIXME: investigate if move selection before compute # speeds up IO. position = position[sel] weight = weight[sel] value = value[sel] else: # workaround a potential dask issue on empty dask arrays position = numpy.empty((0, 3), dtype=Position.dtype) weight = None value = None selection = None if weight is None: weight = numpy.ones(len(position)) if value is None: value = numpy.ones(len(position)) # track total (selected) number and sum of weights Nlocal += len(position) Wlocal += weight.sum() # no interlacing if not self.interlaced: lay = pm.decompose(position, smoothing=0.5 * paintbrush.support) p = lay.exchange(position) w = lay.exchange(weight) v = lay.exchange(value) pm.paint(p, mass=w * v, resampler=paintbrush, hold=True, out=toret) # interlacing: use 2 meshes separated by 1/2 cell size else: lay = pm.decompose(position, smoothing=1.0 * paintbrush.support) p = lay.exchange(position) w = lay.exchange(weight) v = lay.exchange(value) H = pm.BoxSize / pm.Nmesh # in mesh units shifted = pm.affine.shift(0.5) # paint to two shifted meshes pm.paint(p, mass=w * v, resampler=paintbrush, hold=True, out=real1) pm.paint(p, mass=w * v, resampler=paintbrush, transform=shifted, hold=True, out=real2) Nglobal = pm.comm.allreduce(Nlocal) if pm.comm.rank == 0: self.logger.info("painted %d out of %d objects to mesh" % (Nglobal, self.source.csize)) # now the loop over particles is done if not self.interlaced: # nothing to do, toret is already filled. pass else: # compose the two interlaced fields into the final result. c1 = real1.r2c() c2 = real2.r2c() # and then combine for k, s1, s2 in zip(c1.slabs.x, c1.slabs, c2.slabs): kH = sum(k[i] * H[i] for i in range(3)) s1[...] = s1[...] * 0.5 + s2[...] * 0.5 * numpy.exp( 0.5 * 1j * kH) # FFT back to real-space # NOTE: cannot use "toret" here in case user supplied "out" c1.c2r(real1) # need to add to the returned mesh if user supplied "out" toret[:] += real1[:] # unweighted number of objects N = pm.comm.allreduce(Nlocal) # weighted number of objects W = pm.comm.allreduce(Wlocal) # weighted number density (objs/cell) nbar = 1. * W / numpy.prod(pm.Nmesh) # make sure we painted something or nbar is nan; in which case # we set the density to uniform everywhere. if N == 0: warnings.warn(("trying to paint particle source to mesh, " "but no particles were found!"), RuntimeWarning) # shot noise is volume / un-weighted number shotnoise = numpy.prod(pm.BoxSize) / N # save some meta-data toret.attrs = {} toret.attrs['shotnoise'] = shotnoise toret.attrs['N'] = N toret.attrs['W'] = W toret.attrs['num_per_cell'] = nbar csum = toret.csum() if pm.comm.rank == 0: self.logger.info("painted %d out of %d objects to mesh" % (N, self.source.csize)) self.logger.info("mean particles per cell is %g", nbar) self.logger.info("sum is %g ", csum) self.logger.info("normalized the convention to 1 + delta") if normalize: if nbar > 0: toret[...] /= nbar else: toret[...] = 1 return toret