diff --git a/nbodykit/algorithms/convpower.py b/nbodykit/algorithms/convpower.py index 889838615..48599be7f 100644 --- a/nbodykit/algorithms/convpower.py +++ b/nbodykit/algorithms/convpower.py @@ -144,9 +144,9 @@ def __init__(self, first, poles, use_fkp_weights=False, P0_FKP=None): - first = _cast_source(first, Nmesh=Nmesh) + first = _cast_mesh(first, Nmesh=Nmesh) if second is not None: - second = _cast_source(second, Nmesh=Nmesh) + second = _cast_mesh(second, Nmesh=Nmesh) else: second = first @@ -194,22 +194,22 @@ def __init__(self, first, poles, # add FKP weights if use_fkp_weights: - for source in [self.first, self.second]: + for mesh in [self.first, self.second]: if self.comm.rank == 0: - args = (source.fkp_weight, P0_FKP) + args = (mesh.fkp_weight, P0_FKP) self.logger.info("adding FKP weights as the '%s' column, using P0 = %.4e" %args) - for name in ['data', 'randoms']: + for name in ['data', 'randoms']: - # print a warning if we are overwriting a non-default column - old_fkp_weights = source[name][source.fkp_weight] - if source.compute(old_fkp_weights.sum()) != len(old_fkp_weights): - warn = "it appears that we are overwriting FKP weights for the '%s' " %name - warn += "source in FKPCatalog when using 'use_fkp_weights=True' in ConvolvedFFTPower" - warnings.warn(warn) + # print a warning if we are overwriting a non-default column + old_fkp_weights = mesh.source[name][mesh.fkp_weight] + if mesh.source.compute(old_fkp_weights.sum()) != len(old_fkp_weights): + warn = "it appears that we are overwriting FKP weights for the '%s' " %name + warn += "source in FKPCatalog when using 'use_fkp_weights=True' in ConvolvedFFTPower" + warnings.warn(warn) - nbar = source[name][source.nbar] - source[name][source.fkp_weight] = 1.0 / (1. + P0_FKP * nbar) + nbar = mesh.source[name][mesh.nbar] + mesh.source[name][mesh.fkp_weight] = 1.0 / (1. + P0_FKP * nbar) # store meta-data self.attrs = {} @@ -461,7 +461,7 @@ class uses the spherical harmonic addition theorem such that Ylms = [[get_real_Ylm(l,m) for m in range(-l, l+1)] for l in poles[1:]] # paint the 1st FKP density field to the mesh (paints: data - alpha*randoms, essentially) - rfield1 = self.first.paint(Nmesh=self.attrs['Nmesh']) + rfield1 = self.first.compute(Nmesh=self.attrs['Nmesh']) meta1 = rfield1.attrs.copy() if rank == 0: self.logger.info("%s painting of 'first' done" %self.first.window) @@ -485,7 +485,7 @@ class uses the spherical harmonic addition theorem such that if self.first is not self.second: # paint the second field - rfield2 = self.second.paint(Nmesh=self.attrs['Nmesh']) + rfield2 = self.second.compute(Nmesh=self.attrs['Nmesh']) meta2 = rfield2.attrs.copy() if rank == 0: self.logger.info("%s painting of 'second' done" %self.second.window) @@ -670,11 +670,11 @@ def normalization(self, name, alpha): if name+'.norm' not in self.attrs: # the selection (same for first/second) - sel = self.first.compute(self.first[name][self.first.selection]) + sel = self.first.source.compute(self.first.source[name][self.first.selection]) # selected first/second meshes for "name" (data or randoms) - first = self.first[name][sel] - second = self.second[name][sel] + first = self.first.source[name][sel] + second = self.second.source[name][sel] # these are assumed the same b/w first and second meshes comp_weight = first[self.first.comp_weight] @@ -690,7 +690,7 @@ def normalization(self, name, alpha): A = nbar*comp_weight*fkp_weight1*fkp_weight2 if name == 'randoms': A *= alpha - A = self.first.compute(A.sum()) + A = first.compute(A.sum()) self.attrs[name+'.norm'] = self.comm.allreduce(A) return self.attrs[name+'.norm'] @@ -718,11 +718,11 @@ def shotnoise(self, alpha): for name in ['data', 'randoms']: # the selection (same for first/second) - sel = self.first.compute(self.first[name][self.first.selection]) + sel = self.first.source.compute(self.first.source[name][self.first.selection]) # selected first/second meshes for "name" (data or randoms) - first = self.first[name][sel] - second = self.second[name][sel] + first = self.first.source[name][sel] + second = self.second.source[name][sel] # completeness weights (assumed same for first/second) comp_weight = first[self.first.comp_weight] @@ -745,26 +745,26 @@ def shotnoise(self, alpha): # divide by normalization from randoms return Pshot / self.attrs['randoms.norm'] -def _cast_source(source, Nmesh): +def _cast_mesh(mesh, Nmesh): """ Cast an object to a MeshSource. Nmesh is used only on FKPCatalog """ from nbodykit.source.catalog import FKPCatalog from nbodykit.source.catalogmesh import FKPCatalogMesh - if not isinstance(source, (FKPCatalogMesh, FKPCatalog)): + if not isinstance(mesh, (FKPCatalogMesh, FKPCatalog)): raise TypeError("input sources should be a FKPCatalog or FKPCatalogMesh") - if isinstance(source, FKPCatalog): + if isinstance(mesh, FKPCatalog): # if input is CatalogSource, use defaults to make it into a mesh - if not isinstance(source, FKPCatalogMesh): - source = source.to_mesh(Nmesh=Nmesh, dtype='f8', compensated=False) + if not isinstance(mesh, FKPCatalogMesh): + mesh = mesh.to_mesh(Nmesh=Nmesh, dtype='f8', compensated=False) - if Nmesh is not None and any(source.attrs['Nmesh'] != Nmesh): - raise ValueError(("Mismatched Nmesh between __init__ and source.attrs; " + if Nmesh is not None and any(mesh.attrs['Nmesh'] != Nmesh): + raise ValueError(("Mismatched Nmesh between __init__ and mesh.attrs; " "if trying to re-sample with a different mesh, specify " "`Nmesh` as keyword of to_mesh()")) - return source + return mesh def get_compensation(mesh): toret = None @@ -784,7 +784,7 @@ def copy_meta(attrs, meta, prefix=""): def is_valid_crosscorr(first, second): - if second.base is not first.base: + if second.source is not first.source: return False same_cols = ['selection', 'comp_weight', 'nbar'] diff --git a/nbodykit/algorithms/fftpower.py b/nbodykit/algorithms/fftpower.py index e2479d1bf..67c63237a 100644 --- a/nbodykit/algorithms/fftpower.py +++ b/nbodykit/algorithms/fftpower.py @@ -97,13 +97,13 @@ def _compute_3d_power(self): p3d : array_like (complex) the 3D complex array holding the power spectrum """ - c1 = self.first.paint(mode='complex', Nmesh=self.attrs['Nmesh']) + c1 = self.first.compute(mode='complex', Nmesh=self.attrs['Nmesh']) # compute the auto power of single supplied field if self.first is self.second: c2 = c1 else: - c2 = self.second.paint(mode='complex', Nmesh=self.attrs['Nmesh']) + c2 = self.second.compute(mode='complex', Nmesh=self.attrs['Nmesh']) # calculate the 3d power spectrum, slab-by-slab to save memory p3d = c1 @@ -356,13 +356,13 @@ def _compute_3d_power(self): p3d : array_like (complex) the 3D complex array holding the power spectrum """ - c1 = self.first.paint(mode='complex', Nmesh=self.attrs['Nmesh']) + c1 = self.first.compute(mode='complex', Nmesh=self.attrs['Nmesh']) # compute the auto power of single supplied field if self.first is self.second: c2 = c1 else: - c2 = self.second.paint(mode='complex', Nmesh=self.attrs['Nmesh']) + c2 = self.second.compute(mode='complex', Nmesh=self.attrs['Nmesh']) # calculate the 3d power spectrum, slab-by-slab to save memory p3d = c1 @@ -476,7 +476,7 @@ def run(self): - modes : the number of Fourier modes averaged together in each bin """ - c1 = self.first.paint(Nmesh=self.attrs['Nmesh'], mode='complex') + c1 = self.first.compute(Nmesh=self.attrs['Nmesh'], mode='complex') r1 = c1.preview(self.attrs['Nmesh'], axes=self.attrs['axes']) # average along projected axes; # part of product is the rfftn vs r2c (for axes) @@ -487,7 +487,7 @@ def run(self): if self.first is self.second: c2 = c1 else: - c2 = self.second.paint(Nmesh=self.attrs['Nmesh'], mode='complex') + c2 = self.second.compute(Nmesh=self.attrs['Nmesh'], mode='complex') r2 = c2.preview(self.attrs['Nmesh'], axes=self.attrs['axes']) c2 = numpy.fft.rfftn(r2) / self.attrs['Nmesh'].prod() # average along projected axes diff --git a/nbodykit/algorithms/fftrecon.py b/nbodykit/algorithms/fftrecon.py index 671b1d310..406096eb0 100644 --- a/nbodykit/algorithms/fftrecon.py +++ b/nbodykit/algorithms/fftrecon.py @@ -132,7 +132,7 @@ def to_real_field(self): def run(self): s_d, s_r = self._compute_s() - return self._paint(s_d, s_r) + return self._helper_paint(s_d, s_r) def work_with(self, cat, s): pm = self.pm @@ -167,7 +167,7 @@ def _summary_field(self, field, name): self.logger.info("painted %s, mean=%g" % (name, cmean)) - def _paint(self, s_d, s_r): + def _helper_paint(self, s_d, s_r): """ Convert the displacements of data and random to a single reconstruction mesh object. """ def LGS(delta_s_r): diff --git a/nbodykit/base/catalogmesh.py b/nbodykit/base/catalogmesh.py index e24913979..98bbc62f3 100644 --- a/nbodykit/base/catalogmesh.py +++ b/nbodykit/base/catalogmesh.py @@ -9,7 +9,7 @@ from pmesh import window from pmesh.pm import RealField, ComplexField -class CatalogMesh(CatalogSource, MeshSource): +class CatalogMesh(MeshSource): """ A view of a CatalogSource object which knows how to create a MeshSource object from itself. @@ -52,179 +52,36 @@ class CatalogMesh(CatalogSource, MeshSource): logger = logging.getLogger('CatalogMesh') def __repr__(self): - if isinstance(self.base, CatalogMesh): - return repr(self.base) - else: - return "(%s as CatalogMesh)" % repr(self.base) + return "(%s as CatalogMesh)" % repr(self.source) - def __new__(cls, source, BoxSize, Nmesh, dtype, weight, + def __init__(self, source, BoxSize, Nmesh, dtype, weight, value, selection, position='Position', interlaced=False, compensated=False, window='cic', **kwargs): # source here must be a CatalogSource assert isinstance(source, CatalogSourceBase) - # new, empty CatalogSource - obj = CatalogSourceBase.__new__(cls, source.comm) - - # copy over size from the CatalogSource - obj._size = source.size - obj._csize = source.csize - # copy over the necessary meta-data to attrs - obj.attrs['BoxSize'] = BoxSize - obj.attrs['Nmesh'] = Nmesh - obj.attrs['interlaced'] = interlaced - obj.attrs['compensated'] = compensated - obj.attrs['window'] = window + self.attrs['BoxSize'] = BoxSize + self.attrs['Nmesh'] = Nmesh + self.attrs['interlaced'] = interlaced + self.attrs['compensated'] = compensated + self.attrs['window'] = window # copy meta-data from source too - obj.attrs.update(source.attrs) + self.attrs.update(source.attrs) + + self.source = source # store others as straight attributes - obj.dtype = dtype - obj.weight = weight - obj.value = value - obj.selection = selection - obj.position = position + self.dtype = dtype + self.weight = weight + self.value = value + self.selection = selection + self.position = position # add in the Mesh Source attributes - MeshSource.__init__(obj, obj.comm, Nmesh, BoxSize, dtype) - - # finally set the base as the input CatalogSource - # NOTE: set this AFTER MeshSource.__init__() - obj.base = source - - return obj - - def gslice(self, start, stop, end=1, redistribute=True): - """ - Execute a global slice of a CatalogMesh. - - .. note:: - After the global slice is performed, the data is scattered - evenly across all ranks. - - As CatalogMesh objects are views of a CatalogSource, this simply - globally slices the underlying CatalogSource. - - Parameters - ---------- - start : int - the start index of the global slice - stop : int - the stop index of the global slice - step : int, optional - the default step size of the global size - redistribute : bool, optional - if ``True``, evenly re-distribute the sliced data across all - ranks, otherwise just return any local data part of the global - slice - """ - # sort the base object - newbase = self.base.gslice(start, stop, end=end, redistribute=redistribute) - - # view this base class as a CatalogMesh (with default CatalogMesh parameters) - toret = newbase.view(self.__class__) - - # attach the meta-data from self to returned sliced CatalogMesh - return toret.__finalize__(self) - - def sort(self, keys, reverse=False, usecols=None): - """ - Sort the CatalogMesh object globally across all MPI ranks - in ascending order by the input keys. - - Sort columns must be floating or integer type. - - As CatalogMesh objects are views of a CatalogSource, this simply - sorts the underlying CatalogSource. - - Parameters - ---------- - *keys : - the names of columns to sort by. If multiple columns are provided, - the data is sorted consecutively in the order provided - reverse : bool, optional - if ``True``, perform descending sort operations - usecols : list, optional - the name of the columns to include in the returned CatalogSource - """ - # sort the base object - newbase = self.base.sort(keys, reverse=reverse, usecols=usecols) - - # view this base class as a CatalogMesh (with default CatalogMesh parameters) - toret = newbase.view(self.__class__) - - # attach the meta-data from self to returned sliced CatalogMesh - return toret.__finalize__(self) - - def __slice__(self, index): - """ - Return a slice of a CatalogMesh object. - - This slices the CatalogSource object stored as the :attr:`base` - attribute, and then views that sliced object as a CatalogMesh. - - Parameters - ---------- - index : array_like - either a dask or numpy boolean array; this determines which - rows are included in the returned object - - Returns - ------- - subset - the particle source with the same meta-data as ``self``, and - with the sliced data arrays - """ - # this slice of the CatalogSource will be the base of the mesh - base = super(CatalogMesh, self).__slice__(index) - - # view this base class as a CatalogMesh (with default CatalogMesh parameters) - toret = base.view(self.__class__) - - # attach the meta-data from self to returned sliced CatalogMesh - return toret.__finalize__(self) - - def copy(self): - """ - Return a shallow copy of ``self``. - - .. note:: - No copy of data is made. - - Returns - ------- - CatalogMesh : - a new CatalogMesh that holds all of the data columns of ``self`` - """ - # copy the base and view it as a CatalogMesh - toret = self.base.copy().view(self.__class__) - - # attach the meta-data from self to returned sliced CatalogMesh - return toret.__finalize__(self) - - def __finalize__(self, other): - """ - Finalize the creation of a CatalogMesh object by copying over - attributes from a second CatalogMesh. - - This also copies over the relevant MeshSource attributes via a - call to :func:`MeshSource.__finalize__`. - - Parameters - ---------- - obj : CatalogMesh - the second CatalogMesh to copy over attributes from - """ - if isinstance(other, CatalogSourceBase): - self = CatalogSourceBase.__finalize__(self, other) - - if isinstance(other, MeshSource): - self = MeshSource.__finalize__(self, other) - - return self + MeshSource.__init__(self, source.comm, Nmesh, BoxSize, dtype) @property def interlaced(self): @@ -307,7 +164,7 @@ def to_real_field(self, out=None, normalize=True): the painted real field; this has a ``attrs`` dict storing meta-data """ # check for 'Position' column - if self.position not in self: + 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) @@ -342,7 +199,7 @@ def to_real_field(self, out=None, normalize=True): # read the necessary data (as dask arrays) columns = [self.position, self.weight, self.value, self.selection] - Position, Weight, Value, Selection = self.read(columns) + Position, Weight, Value, Selection = self.source.read(columns) # ensure the slices are synced, since decomposition is collective Nlocalmax = max(pm.comm.allgather(len(Position))) @@ -356,11 +213,11 @@ def to_real_field(self, out=None, normalize=True): if len(Position) != 0: # selection has to be computed many times when data is `large`. - sel = self.base.compute(Selection[s]) + sel = self.source.compute(Selection[s]) # be sure to use the source to compute position, weight, value = \ - self.base.compute(Position[s], Weight[s], Value[s]) + self.source.compute(Position[s], Weight[s], Value[s]) # FIXME: investigate if move selection before compute # speeds up IO. @@ -412,7 +269,7 @@ def to_real_field(self, out=None, normalize=True): if pm.comm.rank == 0: self.logger.info("painted %d out of %d objects to mesh" - % (Nglobal, self.base.csize)) + % (Nglobal, self.source.csize)) # now the loop over particles is done @@ -466,7 +323,7 @@ def to_real_field(self, out=None, normalize=True): csum = toret.csum() if pm.comm.rank == 0: - self.logger.info("painted %d out of %d objects to mesh" %(N,self.base.csize)) + 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") diff --git a/nbodykit/base/mesh.py b/nbodykit/base/mesh.py index fad3ea1bb..617762474 100644 --- a/nbodykit/base/mesh.py +++ b/nbodykit/base/mesh.py @@ -1,6 +1,7 @@ import numpy import logging from pmesh.pm import ParticleMesh, RealField, BaseComplexField +import warnings class MeshSource(object): """ @@ -230,7 +231,17 @@ def to_field(self, mode='real', out=None): return var + def compute(self, mode='real', Nmesh=None): + """ + Compute / Fetch the mesh object into memory as a RealField or ComplexField object. + """ + return self._paint_XXX(mode=mode, Nmesh=Nmesh) + def paint(self, mode="real", Nmesh=None): + warnings.warn("the paint method is deprecated from the Public API. Use .compute() instead.", DeprecationWarning) + return self._paint_XXX(mode=mode, Nmesh=Nmesh) + + def _paint_XXX(self, mode="real", Nmesh=None): """ Paint the density on the mesh and apply any transformation functions specified in :attr:`actions`. @@ -360,7 +371,7 @@ def save(self, output, dataset='Field', mode='real'): import json from nbodykit.utils import JSONEncoder - field = self.paint(mode=mode) + field = self.compute(mode=mode) with bigfile.BigFileMPI(self.pm.comm, output, create=True) as ff: data = numpy.empty(shape=field.size, dtype=field.dtype) diff --git a/nbodykit/base/tests/test_catalog.py b/nbodykit/base/tests/test_catalog.py index e9e0d7401..9f78b85f1 100644 --- a/nbodykit/base/tests/test_catalog.py +++ b/nbodykit/base/tests/test_catalog.py @@ -78,10 +78,8 @@ def test_tomesh(comm): source['Weight2'] = source['Velocity'][:, 2] mesh = source.to_mesh(Nmesh=128, compensated=True) - assert_allclose(source['Position'], mesh['Position']) mesh = source.to_mesh(Nmesh=128, compensated=True, interlaced=True) - assert_allclose(source['Position'], mesh['Position']) mesh = source.to_mesh(Nmesh=128, weight='Weight0') diff --git a/nbodykit/base/tests/test_catalogmesh.py b/nbodykit/base/tests/test_catalogmesh.py index 4b49625dc..9636c9cdd 100644 --- a/nbodykit/base/tests/test_catalogmesh.py +++ b/nbodykit/base/tests/test_catalogmesh.py @@ -8,64 +8,11 @@ # debug logging setup_logging("debug") - -@MPITest([1, 4]) -def test_gslice(comm): - CurrentMPIComm.set(comm) - - # the catalog to slice - d = UniformCatalog(1000, 1.0, seed=42) - d = d.to_mesh(Nmesh=32) - - sels = [(0,10,1), (None,10,1), (0,50,4), (0,50,-1), (-10,-1,1), (-10,None,1)] - for (start,stop,end) in sels: - - sliced = d.gslice(start,stop,end) - - for col in d: - data1 = numpy.concatenate(comm.allgather(d[col]), axis=0) - data2 = numpy.concatenate(comm.allgather(sliced[col]), axis=0) - - sl = slice(start,stop,end) - assert_array_equal(data1[sl], data2, err_msg="using slice= "+str(sl)) - - # empty slice - sliced = d.gslice(0,0) - assert len(sliced) == 0 - -@MPITest([1, 4]) -def test_sort_ascending(comm): - CurrentMPIComm.set(comm) - - # the mesh to sort - d = UniformCatalog(100, 1.0) - d['mass'] = 10**(d.rng.uniform(low=12, high=15)) - mesh = d.to_mesh(Nmesh=32) - - # invalid sort key - with pytest.raises(ValueError): - mesh2 = mesh.sort('BadColumn', reverse=False) - - # duplicate sort keys - with pytest.raises(ValueError): - mesh2 = mesh.sort(['mass', 'mass']) - - # sort in ascending order by mass - mesh2 = mesh.sort('mass', reverse=False) - - # make sure we have all the columns - assert all(col in mesh2 for col in mesh) - - mass = numpy.concatenate(comm.allgather(mesh['mass'])) - sorted_mass = numpy.concatenate(comm.allgather(mesh2['mass'])) - assert_array_equal(numpy.sort(mass), sorted_mass) - - @MPITest([4]) def test_tsc_interlacing(comm): CurrentMPIComm.set(comm) - source = UniformCatalog(nbar=3e-2, BoxSize=512., seed=42) + source = UniformCatalog(nbar=3e-4, BoxSize=512., seed=42) # interlacing with TSC mesh = source.to_mesh(window='tsc', Nmesh=64, interlaced=True, compensated=True) @@ -78,16 +25,16 @@ def test_tsc_interlacing(comm): def test_paint_chunksize(comm): CurrentMPIComm.set(comm) - source = UniformCatalog(nbar=3e-2, BoxSize=512., seed=42) + source = UniformCatalog(nbar=3e-4, BoxSize=512., seed=42) # interlacing with TSC mesh = source.to_mesh(window='tsc', Nmesh=64, interlaced=True, compensated=True) with set_options(paint_chunk_size=source.csize // 4): - r1 = mesh.paint() + r1 = mesh.compute() with set_options(paint_chunk_size=source.csize): - r2 = mesh.paint() + r2 = mesh.compute() assert_allclose(r1, r2) @@ -95,7 +42,7 @@ def test_paint_chunksize(comm): def test_cic_interlacing(comm): CurrentMPIComm.set(comm) - source = UniformCatalog(nbar=3e-2, BoxSize=512., seed=42) + source = UniformCatalog(nbar=3e-4, BoxSize=512., seed=42) # interlacing with TSC mesh = source.to_mesh(window='cic', Nmesh=64, interlaced=True, compensated=True) @@ -108,7 +55,7 @@ def test_cic_interlacing(comm): def test_setters(comm): CurrentMPIComm.set(comm) - source = UniformCatalog(nbar=3e-2, BoxSize=512., seed=42) + source = UniformCatalog(nbar=3e-4, BoxSize=512., seed=42) # make the mesh mesh = source.to_mesh(window='cic', Nmesh=64, interlaced=True, compensated=True) @@ -129,7 +76,7 @@ def test_setters(comm): def test_bad_window(comm): CurrentMPIComm.set(comm) - source = UniformCatalog(nbar=3e-2, BoxSize=512., seed=42) + source = UniformCatalog(nbar=3e-4, BoxSize=512., seed=42) # make the mesh mesh = source.to_mesh(window='cic', Nmesh=64, interlaced=True, compensated=True) @@ -142,7 +89,7 @@ def test_bad_window(comm): def test_no_compensation(comm): CurrentMPIComm.set(comm) - source = UniformCatalog(nbar=3e-2, BoxSize=512., seed=42) + source = UniformCatalog(nbar=3e-4, BoxSize=512., seed=42) # make the mesh mesh = source.to_mesh(window='cic', Nmesh=64, interlaced=True, compensated=True) @@ -154,57 +101,6 @@ def test_no_compensation(comm): with pytest.raises(ValueError): actions = mesh.actions -@MPITest([1, 4]) -def test_copy(comm): - - CurrentMPIComm.set(comm) - source = UniformCatalog(nbar=0.2e-3, BoxSize=1024., seed=42) - source['TEST'] = 10 - source.attrs['TEST'] = 'TEST' - - # make the mesh - source = source.to_mesh(Nmesh=32) - - # store original data - data = {} - for col in source: - data[col] = source[col].compute() - - # make copy - copy = source.copy() - - # modify original - source['Position'] += 100. - source['Velocity'] *= 10. - - # check data is equal to original - for col in copy: - assert_array_equal(copy[col].compute(), data[col]) - - # check meta-data - for k in source.attrs: - assert k in copy.attrs - -@MPITest([4]) -def test_slice(comm): - CurrentMPIComm.set(comm) - - # the CatalogSource - source = UniformCatalog(nbar=0.2e-3, BoxSize=1024., seed=42) - source['TEST'] = 10. - - # the mesh - source = source.to_mesh(Nmesh=32) - - # slice a subset - subset = source[:10] - assert all(col in subset for col in source.columns) - assert isinstance(subset, source.__class__) - assert len(subset) == 10 - assert_array_equal(subset['Position'], source['Position'].compute()[:10]) - - subset = source[[0,1,2]] - assert_array_equal(subset['Position'], source['Position'].compute()[[0,1,2]]) @MPITest([4]) def test_view(comm): @@ -221,17 +117,13 @@ def test_view(comm): # view view = mesh.view() assert view.base is mesh - assert isinstance(view, mesh.__class__) + from nbodykit.base.mesh import MeshSource + assert isinstance(view, MeshSource) # check meta-data for k in mesh.attrs: assert k in view.attrs - # adding columns to the view changes original source - view['TEST2'] = 5.0 - assert 'TEST2' in source -class CodeReached(BaseException): pass - @MPITest([1, 4]) def test_apply_nocompensation(comm): CurrentMPIComm.set(comm) @@ -251,7 +143,7 @@ def raisefunc(k, v): mesh = mesh.apply(raisefunc) with pytest.raises(StopIteration): - mesh.paint() + mesh.compute() # view view = mesh.view() @@ -281,5 +173,5 @@ def raisefunc(k, v): mesh = mesh.apply(raisefunc) with pytest.raises(StopIteration): - mesh.paint() + mesh.compute() diff --git a/nbodykit/base/tests/test_mesh.py b/nbodykit/base/tests/test_mesh.py index bf9a07c7a..31f061208 100644 --- a/nbodykit/base/tests/test_mesh.py +++ b/nbodykit/base/tests/test_mesh.py @@ -69,7 +69,7 @@ def test_real_save(comm): assert_array_equal(source2.attrs[k], source.attrs[k]) # check data - assert_array_equal(source2.paint(mode='complex'), source.paint(mode='complex')) + assert_array_equal(source2.compute(mode='complex'), source.compute(mode='complex')) # cleanup comm.barrier() @@ -107,7 +107,7 @@ def test_real_save(comm): assert_array_equal(source2.attrs[k], source.attrs[k]) # check data - assert_array_equal(source2.paint(mode='real'), source.paint(mode='real')) + assert_array_equal(source2.compute(mode='real'), source.compute(mode='real')) # cleanup comm.barrier() @@ -125,7 +125,7 @@ def test_preview(comm): source = LinearMesh(Plin, Nmesh=64, BoxSize=512, seed=42) # the painted RealField - real = source.paint(mode='real') + real = source.compute(mode='real') preview = source.preview() assert_allclose(preview.sum(), real.csum(), rtol=1e-5) @@ -145,7 +145,7 @@ def test_resample(comm): source = LinearMesh(Plin, Nmesh=64, BoxSize=512, seed=42) # re-sample to Nmesh=32 - real = source.paint(mode='real', Nmesh=32) + real = source.compute(mode='real', Nmesh=32) # and preview at same resolution preview = source.preview(Nmesh=32) @@ -172,7 +172,7 @@ def test_bad_mode(comm): field = source.to_field(mode='BAD') with pytest.raises(ValueError): - field = source.paint(mode='BAD') + field = source.compute(mode='BAD') @MPITest([4]) def test_view(comm): diff --git a/nbodykit/source/catalog/tests/test_lognormal.py b/nbodykit/source/catalog/tests/test_lognormal.py index 3d4f157db..cba414d37 100644 --- a/nbodykit/source/catalog/tests/test_lognormal.py +++ b/nbodykit/source/catalog/tests/test_lognormal.py @@ -20,7 +20,7 @@ def test_lognormal_sparse(comm): mesh = source.to_mesh(compensated=False) - real = mesh.paint(mode='real') + real = mesh.compute(mode='real') assert_allclose(real.cmean(), 1.0) @MPITest([1, 4]) @@ -32,7 +32,7 @@ def test_lognormal_dense(comm): source = LogNormalCatalog(Plin=Plin, nbar=0.2e-2, BoxSize=128., Nmesh=8, seed=42) mesh = source.to_mesh(compensated=False) - real = mesh.paint(mode='real') + real = mesh.compute(mode='real') assert_allclose(real.cmean(), 1.0, rtol=1e-5) @MPITest([4]) @@ -62,7 +62,7 @@ def test_lognormal_velocity(comm): source['Value'] = source['Velocity'][:, 0]**2 mesh = source.to_mesh(compensated=False) - real = mesh.paint(mode='real') + real = mesh.compute(mode='real') velsum = comm.allreduce((source['Velocity'][:, 0]**2).sum().compute()) velmean = velsum / source.csize diff --git a/nbodykit/source/catalogmesh/fkp.py b/nbodykit/source/catalogmesh/fkp.py index 27e5c94d5..c3d5d6163 100644 --- a/nbodykit/source/catalogmesh/fkp.py +++ b/nbodykit/source/catalogmesh/fkp.py @@ -42,7 +42,7 @@ class FKPCatalogMesh(MultipleSpeciesCatalogMesh): """ logger = logging.getLogger('FKPCatalogMesh') - def __new__(cls, source, BoxSize, Nmesh, dtype, selection, + def __init__(self, source, BoxSize, Nmesh, dtype, selection, comp_weight, fkp_weight, nbar, value='Value', position='Position', interlaced=False, compensated=False, window='cic'): @@ -55,16 +55,14 @@ def __new__(cls, source, BoxSize, Nmesh, dtype, selection, position = '_RecenteredPosition' weight = '_TotalWeight' - obj = super(FKPCatalogMesh, cls).__new__(cls, source, BoxSize, Nmesh, + MultipleSpeciesCatalogMesh.__init__(self, source, BoxSize, Nmesh, dtype, weight, value, selection, position=position, interlaced=interlaced, compensated=compensated, window=window) - obj._uncentered_position = uncentered_position - obj.comp_weight = comp_weight - obj.fkp_weight = fkp_weight - obj.nbar = nbar - - return obj + self._uncentered_position = uncentered_position + self.comp_weight = comp_weight + self.fkp_weight = fkp_weight + self.nbar = nbar def recenter_box(self, BoxSize, BoxCenter): """ @@ -82,7 +80,7 @@ def recenter_box(self, BoxSize, BoxCenter): # update meta-data for val, name in zip([BoxSize, BoxCenter], ['BoxSize', 'BoxCenter']): self.attrs[name] = val - self.base.attrs[name] = val + self.source.attrs[name] = val def to_real_field(self): @@ -127,18 +125,18 @@ def to_real_field(self): the field object holding the FKP density field in real space """ # add necessary FKP columns for INTERNAL use - for name in self.base.species: + for name in self.source.species: # a total weight for the mesh is completeness weight x FKP weight - self[name]['_TotalWeight'] = self.TotalWeight(name) + self.source[name]['_TotalWeight'] = self.TotalWeight(name) # position on the mesh is re-centered to [-BoxSize/2, BoxSize/2] - self[name]['_RecenteredPosition'] = self.RecenteredPosition(name) + self.source[name]['_RecenteredPosition'] = self.RecenteredPosition(name) attrs = {} # determine alpha, the weighted number ratio - for name in self.base.species: + for name in self.source.species: attrs[name+'.W'] = self.weighted_total(name) attrs['alpha'] = attrs['data.W'] / attrs['randoms.W'] @@ -164,9 +162,9 @@ def to_real_field(self): real.attrs.pop('randoms.shotnoise', None) # delete internal columns - for name in self.base.species: - del self[name+'/_RecenteredPosition'] - del self[name+'/_TotalWeight'] + for name in self.source.species: + del self.source[name+'/_RecenteredPosition'] + del self.source[name+'/_TotalWeight'] return real @@ -179,7 +177,7 @@ def RecenteredPosition(self, name): position array. """ assert name in ['data', 'randoms'] - return self[name][self._uncentered_position] - self.attrs['BoxCenter'] + return self.source[name][self._uncentered_position] - self.attrs['BoxCenter'] def TotalWeight(self, name): """ @@ -187,7 +185,7 @@ def TotalWeight(self, name): the FKP weight. """ assert name in ['data', 'randoms'] - return self[name][self.comp_weight] * self[name][self.fkp_weight] + return self.source[name][self.comp_weight] * self.source[name][self.fkp_weight] def weighted_total(self, name): r""" @@ -201,11 +199,11 @@ def weighted_total(self, name): W = \sum w_\mathrm{comp} """ # the selection - sel = self.compute(self[name][self.selection]) + sel = self.source.compute(self.source[name][self.selection]) # the selected mesh for "name" - mesh = self[name][sel] + selected = self.source[name][sel] # sum up completeness weights - wsum = self.compute(mesh[self.comp_weight].sum()) + wsum = self.source.compute(selected[self.comp_weight].sum()) return self.comm.allreduce(wsum) diff --git a/nbodykit/source/catalogmesh/species.py b/nbodykit/source/catalogmesh/species.py index e2f4493ab..70b587785 100644 --- a/nbodykit/source/catalogmesh/species.py +++ b/nbodykit/source/catalogmesh/species.py @@ -12,7 +12,7 @@ class MultipleSpeciesCatalogMesh(CatalogMesh): designed to paint the density field from a sum of multiple types of particles. - The :func:`paint` function paints the density field summed over + The :func:`compute` function paints the density field summed over all particle species. Parameters @@ -39,14 +39,17 @@ class MultipleSpeciesCatalogMesh(CatalogMesh): """ logger = logging.getLogger('MultipleSpeciesCatalogMesh') - def __new__(cls, source, *args, **kwargs): + def __init__(self, source, *args, **kwargs): from nbodykit.source.catalog import MultipleSpeciesCatalog if not isinstance(source, MultipleSpeciesCatalog): raise TypeError(("the input source for MultipleSpeciesCatalogMesh " "must be a MultipleSpeciesCatalog")) - return super(MultipleSpeciesCatalogMesh, cls).__new__(cls, source, *args, **kwargs) + CatalogMesh.__init__(self, source, *args, **kwargs) + + def __iter__(self): + return iter(self.source) def __getitem__(self, key): """ @@ -58,19 +61,29 @@ def __getitem__(self, key): :func:`CatalogSource.__getitem__`. """ # return a Mesh holding only the specific species - if isinstance(key, string_types) and key in self.base.species: + if isinstance(key, string_types) and key in self.source.species: # CatalogSource holding only requested species - cat = self.base[key] + cat = self.source[key] # view as a catalog mesh - mesh = cat.view(CatalogMesh) + mesh = CatalogMesh(cat, + BoxSize=self.attrs['BoxSize'], + Nmesh=self.attrs['Nmesh'], + dtype=self.dtype, + weight=self.weight, + value=self.value, + selection=self.selection, + position=self.position, + interlaced=self.interlaced, + compensated=self.compensated, + window=self.window, + ) # attach attributes from self return mesh.__finalize__(self) - - # return the base class behavior - return CatalogMesh.__getitem__(self, key) + else: + raise KeyError("species '%s' not found" % key) def to_real_field(self, normalize=True): r""" @@ -110,7 +123,7 @@ def to_real_field(self, normalize=True): real = self.pm.create(mode='real', value=0) # loop over each species - for name in self.base.species: + for name in self.source.species: if self.pm.comm.rank == 0: self.logger.info("painting the '%s' species" %name) @@ -138,8 +151,8 @@ def to_real_field(self, normalize=True): # compute total shot noise by summing of shot noise each species real.attrs['shotnoise'] = 0 - total_weight = sum(real.attrs['%s.W' %name] for name in self.base.species) - for name in self.base.species: + total_weight = sum(real.attrs['%s.W' %name] for name in self.source.species) + for name in self.source.species: this_Pshot = real.attrs['%s.shotnoise' %name] this_weight = real.attrs['%s.W' %name] real.attrs['shotnoise'] += (this_weight/total_weight)**2 * this_Pshot diff --git a/nbodykit/source/catalogmesh/tests/test_fkp.py b/nbodykit/source/catalogmesh/tests/test_fkp.py index 4413dc4ec..61e6e2564 100644 --- a/nbodykit/source/catalogmesh/tests/test_fkp.py +++ b/nbodykit/source/catalogmesh/tests/test_fkp.py @@ -43,12 +43,12 @@ def test_paint(comm): mesh2 = source2.to_mesh(Nmesh=32, BoxSize=512) # update weights for source1 and source2 - mesh1['Weight'] *= mesh1['FKPWeight'] - mesh2['Weight'] *= mesh2['FKPWeight'] + mesh1.source['Weight'] *= mesh1.source['FKPWeight'] + mesh2.source['Weight'] *= mesh2.source['FKPWeight'] # paint the re-centered Position - mesh1['Position'] -= mesh.attrs['BoxCenter'] - mesh2['Position'] -= mesh.attrs['BoxCenter'] + mesh1.source['Position'] -= mesh.attrs['BoxCenter'] + mesh2.source['Position'] -= mesh.attrs['BoxCenter'] # alpha is the sum of Weight alpha = 1. * source1.csize * WEIGHT1 / (source2.csize * WEIGHT2) diff --git a/nbodykit/source/catalogmesh/tests/test_species.py b/nbodykit/source/catalogmesh/tests/test_species.py index c539683d9..e9bc587bf 100644 --- a/nbodykit/source/catalogmesh/tests/test_species.py +++ b/nbodykit/source/catalogmesh/tests/test_species.py @@ -44,12 +44,10 @@ def test_getitem(comm): for source, name in zip([source1, source2], ['data', 'randoms']): submesh = mesh[name] # should be equal to source - for col in source: - assert col in submesh - assert_array_equal(submesh[col].compute(), source[col].compute()) + assert submesh.source is cat[name] @MPITest([1, 4]) -def test_paint(comm): +def test_compute(comm): CurrentMPIComm.set(comm) @@ -111,7 +109,7 @@ def test_paint_interlaced(comm): # the combined density field #combined = mesh.to_real_field() - combined = mesh.paint() + combined = mesh.compute() assert_allclose(combined.cmean(), 1.0) # must be the same diff --git a/nbodykit/source/mesh/tests/test_bigfile.py b/nbodykit/source/mesh/tests/test_bigfile.py index 2c4695432..c92bf2826 100644 --- a/nbodykit/source/mesh/tests/test_bigfile.py +++ b/nbodykit/source/mesh/tests/test_bigfile.py @@ -19,8 +19,8 @@ def test_bigfile_grid(comm): Plin = cosmology.LinearPower(cosmo, redshift=0.55, transfer='EisensteinHu') source = LinearMesh(Plin, BoxSize=512, Nmesh=64, seed=42) - real = source.paint(mode='real') - complex = source.paint(mode="complex") + real = source.compute(mode='real') + complex = source.compute(mode="complex") # and save to tmp directory if comm.rank == 0: @@ -33,7 +33,7 @@ def test_bigfile_grid(comm): # now load it and paint to the algorithm's ParticleMesh source = BigFileMesh(path=output, dataset='Field') - loaded_real = source.paint() + loaded_real = source.compute() # compare to direct algorithm result assert_array_equal(real, loaded_real) @@ -42,7 +42,7 @@ def test_bigfile_grid(comm): # now load it and paint to the algorithm's ParticleMesh source = BigFileMesh(path=output, dataset='FieldC') - loaded_real = source.paint(mode="complex") + loaded_real = source.compute(mode="complex") # compare to direct algorithm result assert_allclose(complex, loaded_real, atol=1e-7) diff --git a/nbodykit/tests/test_lab.py b/nbodykit/tests/test_lab.py index e41854e72..85062fa41 100644 --- a/nbodykit/tests/test_lab.py +++ b/nbodykit/tests/test_lab.py @@ -26,7 +26,7 @@ def test_fftpower(comm): result.save(output) @MPITest([1, 4]) -def test_paint(comm): +def test_compute(comm): cosmo = cosmology.Planck15 CurrentMPIComm.set(comm) @@ -48,8 +48,8 @@ def filter(k, v): source = source.apply(filter) - real = source.paint(mode='real') - complex = source.paint(mode='complex') + real = source.compute(mode='real') + complex = source.compute(mode='complex') source.save(output="./test_paint-real-%d.bigfile" % comm.size, mode='real') source.save(output="./test_paint-complex-%d.bigfile" % comm.size, mode='complex')