Skip to content

Commit ef10f2d

Browse files
committed
Remove Catalog interface from CatalogMesh.
This is required to add a compute method to MeshSource. Because if CatalogMesh is a CatalogSource then CatalogSource.compute comes into the way. I think this PR also sheds some light on how to deal with the temporary variables created during FKP/convpower. They probably shall be saved as attributes of the CatalogMesh objects directly, rather than stored as columns of the underlying sources.
1 parent 6f0b167 commit ef10f2d

14 files changed

Lines changed: 145 additions & 298 deletions

File tree

nbodykit/algorithms/convpower.py

Lines changed: 31 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -144,9 +144,9 @@ def __init__(self, first, poles,
144144
use_fkp_weights=False,
145145
P0_FKP=None):
146146

147-
first = _cast_source(first, Nmesh=Nmesh)
147+
first = _cast_mesh(first, Nmesh=Nmesh)
148148
if second is not None:
149-
second = _cast_source(second, Nmesh=Nmesh)
149+
second = _cast_mesh(second, Nmesh=Nmesh)
150150
else:
151151
second = first
152152

@@ -194,22 +194,22 @@ def __init__(self, first, poles,
194194

195195
# add FKP weights
196196
if use_fkp_weights:
197-
for source in [self.first, self.second]:
197+
for mesh in [self.first, self.second]:
198198
if self.comm.rank == 0:
199-
args = (source.fkp_weight, P0_FKP)
199+
args = (mesh.fkp_weight, P0_FKP)
200200
self.logger.info("adding FKP weights as the '%s' column, using P0 = %.4e" %args)
201201

202-
for name in ['data', 'randoms']:
202+
for name in ['data', 'randoms']:
203203

204-
# print a warning if we are overwriting a non-default column
205-
old_fkp_weights = source[name][source.fkp_weight]
206-
if source.compute(old_fkp_weights.sum()) != len(old_fkp_weights):
207-
warn = "it appears that we are overwriting FKP weights for the '%s' " %name
208-
warn += "source in FKPCatalog when using 'use_fkp_weights=True' in ConvolvedFFTPower"
209-
warnings.warn(warn)
204+
# print a warning if we are overwriting a non-default column
205+
old_fkp_weights = mesh.source[name][mesh.fkp_weight]
206+
if mesh.source.compute(old_fkp_weights.sum()) != len(old_fkp_weights):
207+
warn = "it appears that we are overwriting FKP weights for the '%s' " %name
208+
warn += "source in FKPCatalog when using 'use_fkp_weights=True' in ConvolvedFFTPower"
209+
warnings.warn(warn)
210210

211-
nbar = source[name][source.nbar]
212-
source[name][source.fkp_weight] = 1.0 / (1. + P0_FKP * nbar)
211+
nbar = mesh.source[name][mesh.nbar]
212+
mesh.source[name][mesh.fkp_weight] = 1.0 / (1. + P0_FKP * nbar)
213213

214214
# store meta-data
215215
self.attrs = {}
@@ -461,7 +461,7 @@ class uses the spherical harmonic addition theorem such that
461461
Ylms = [[get_real_Ylm(l,m) for m in range(-l, l+1)] for l in poles[1:]]
462462

463463
# paint the 1st FKP density field to the mesh (paints: data - alpha*randoms, essentially)
464-
rfield1 = self.first.paint(Nmesh=self.attrs['Nmesh'])
464+
rfield1 = self.first.compute(Nmesh=self.attrs['Nmesh'])
465465
meta1 = rfield1.attrs.copy()
466466
if rank == 0:
467467
self.logger.info("%s painting of 'first' done" %self.first.window)
@@ -485,7 +485,7 @@ class uses the spherical harmonic addition theorem such that
485485
if self.first is not self.second:
486486

487487
# paint the second field
488-
rfield2 = self.second.paint(Nmesh=self.attrs['Nmesh'])
488+
rfield2 = self.second.compute(Nmesh=self.attrs['Nmesh'])
489489
meta2 = rfield2.attrs.copy()
490490
if rank == 0: self.logger.info("%s painting of 'second' done" %self.second.window)
491491

@@ -670,11 +670,11 @@ def normalization(self, name, alpha):
670670
if name+'.norm' not in self.attrs:
671671

672672
# the selection (same for first/second)
673-
sel = self.first.compute(self.first[name][self.first.selection])
673+
sel = self.first.source.compute(self.first.source[name][self.first.selection])
674674

675675
# selected first/second meshes for "name" (data or randoms)
676-
first = self.first[name][sel]
677-
second = self.second[name][sel]
676+
first = self.first.source[name][sel]
677+
second = self.second.source[name][sel]
678678

679679
# these are assumed the same b/w first and second meshes
680680
comp_weight = first[self.first.comp_weight]
@@ -690,7 +690,7 @@ def normalization(self, name, alpha):
690690
A = nbar*comp_weight*fkp_weight1*fkp_weight2
691691
if name == 'randoms':
692692
A *= alpha
693-
A = self.first.compute(A.sum())
693+
A = first.compute(A.sum())
694694
self.attrs[name+'.norm'] = self.comm.allreduce(A)
695695

696696
return self.attrs[name+'.norm']
@@ -718,11 +718,11 @@ def shotnoise(self, alpha):
718718
for name in ['data', 'randoms']:
719719

720720
# the selection (same for first/second)
721-
sel = self.first.compute(self.first[name][self.first.selection])
721+
sel = self.first.source.compute(self.first.source[name][self.first.selection])
722722

723723
# selected first/second meshes for "name" (data or randoms)
724-
first = self.first[name][sel]
725-
second = self.second[name][sel]
724+
first = self.first.source[name][sel]
725+
second = self.second.source[name][sel]
726726

727727
# completeness weights (assumed same for first/second)
728728
comp_weight = first[self.first.comp_weight]
@@ -745,26 +745,26 @@ def shotnoise(self, alpha):
745745
# divide by normalization from randoms
746746
return Pshot / self.attrs['randoms.norm']
747747

748-
def _cast_source(source, Nmesh):
748+
def _cast_mesh(mesh, Nmesh):
749749
"""
750750
Cast an object to a MeshSource. Nmesh is used only on FKPCatalog
751751
"""
752752
from nbodykit.source.catalog import FKPCatalog
753753
from nbodykit.source.catalogmesh import FKPCatalogMesh
754-
if not isinstance(source, (FKPCatalogMesh, FKPCatalog)):
754+
if not isinstance(mesh, (FKPCatalogMesh, FKPCatalog)):
755755
raise TypeError("input sources should be a FKPCatalog or FKPCatalogMesh")
756756

757-
if isinstance(source, FKPCatalog):
757+
if isinstance(mesh, FKPCatalog):
758758
# if input is CatalogSource, use defaults to make it into a mesh
759-
if not isinstance(source, FKPCatalogMesh):
760-
source = source.to_mesh(Nmesh=Nmesh, dtype='f8', compensated=False)
759+
if not isinstance(mesh, FKPCatalogMesh):
760+
mesh = mesh.to_mesh(Nmesh=Nmesh, dtype='f8', compensated=False)
761761

762-
if Nmesh is not None and any(source.attrs['Nmesh'] != Nmesh):
763-
raise ValueError(("Mismatched Nmesh between __init__ and source.attrs; "
762+
if Nmesh is not None and any(mesh.attrs['Nmesh'] != Nmesh):
763+
raise ValueError(("Mismatched Nmesh between __init__ and mesh.attrs; "
764764
"if trying to re-sample with a different mesh, specify "
765765
"`Nmesh` as keyword of to_mesh()"))
766766

767-
return source
767+
return mesh
768768

769769
def get_compensation(mesh):
770770
toret = None
@@ -784,7 +784,7 @@ def copy_meta(attrs, meta, prefix=""):
784784

785785
def is_valid_crosscorr(first, second):
786786

787-
if second.base is not first.base:
787+
if second.source is not first.source:
788788
return False
789789

790790
same_cols = ['selection', 'comp_weight', 'nbar']

nbodykit/algorithms/fftpower.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -97,13 +97,13 @@ def _compute_3d_power(self):
9797
p3d : array_like (complex)
9898
the 3D complex array holding the power spectrum
9999
"""
100-
c1 = self.first.paint(mode='complex', Nmesh=self.attrs['Nmesh'])
100+
c1 = self.first.compute(mode='complex', Nmesh=self.attrs['Nmesh'])
101101

102102
# compute the auto power of single supplied field
103103
if self.first is self.second:
104104
c2 = c1
105105
else:
106-
c2 = self.second.paint(mode='complex', Nmesh=self.attrs['Nmesh'])
106+
c2 = self.second.compute(mode='complex', Nmesh=self.attrs['Nmesh'])
107107

108108
# calculate the 3d power spectrum, slab-by-slab to save memory
109109
p3d = c1
@@ -356,13 +356,13 @@ def _compute_3d_power(self):
356356
p3d : array_like (complex)
357357
the 3D complex array holding the power spectrum
358358
"""
359-
c1 = self.first.paint(mode='complex', Nmesh=self.attrs['Nmesh'])
359+
c1 = self.first.compute(mode='complex', Nmesh=self.attrs['Nmesh'])
360360

361361
# compute the auto power of single supplied field
362362
if self.first is self.second:
363363
c2 = c1
364364
else:
365-
c2 = self.second.paint(mode='complex', Nmesh=self.attrs['Nmesh'])
365+
c2 = self.second.compute(mode='complex', Nmesh=self.attrs['Nmesh'])
366366

367367
# calculate the 3d power spectrum, slab-by-slab to save memory
368368
p3d = c1
@@ -476,7 +476,7 @@ def run(self):
476476
- modes :
477477
the number of Fourier modes averaged together in each bin
478478
"""
479-
c1 = self.first.paint(Nmesh=self.attrs['Nmesh'], mode='complex')
479+
c1 = self.first.compute(Nmesh=self.attrs['Nmesh'], mode='complex')
480480
r1 = c1.preview(self.attrs['Nmesh'], axes=self.attrs['axes'])
481481
# average along projected axes;
482482
# part of product is the rfftn vs r2c (for axes)
@@ -487,7 +487,7 @@ def run(self):
487487
if self.first is self.second:
488488
c2 = c1
489489
else:
490-
c2 = self.second.paint(Nmesh=self.attrs['Nmesh'], mode='complex')
490+
c2 = self.second.compute(Nmesh=self.attrs['Nmesh'], mode='complex')
491491
r2 = c2.preview(self.attrs['Nmesh'], axes=self.attrs['axes'])
492492
c2 = numpy.fft.rfftn(r2) / self.attrs['Nmesh'].prod() # average along projected axes
493493

nbodykit/algorithms/fftrecon.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ def to_real_field(self):
132132
def run(self):
133133

134134
s_d, s_r = self._compute_s()
135-
return self._paint(s_d, s_r)
135+
return self._helper_paint(s_d, s_r)
136136

137137
def work_with(self, cat, s):
138138
pm = self.pm
@@ -167,7 +167,7 @@ def _summary_field(self, field, name):
167167
self.logger.info("painted %s, mean=%g" % (name, cmean))
168168

169169

170-
def _paint(self, s_d, s_r):
170+
def _helper_paint(self, s_d, s_r):
171171
""" Convert the displacements of data and random to a single reconstruction mesh object. """
172172

173173
def LGS(delta_s_r):

0 commit comments

Comments
 (0)