Interlaced species

nbodykit/base/catalogmesh.py

@@ -2,6 +2,7 @@
 from nbodykit.base.catalog import CatalogSource, CatalogSourceBase
 import numpy
 import logging
+import warnings
 
 # for converting from particle to mesh
 from pmesh import window
@@ -330,12 +331,8 @@ def to_real_field(self, out=None, normalize=True):
         # since out may have non-zero elements, messing up our interlacing sum
         if self.interlaced:
 
-            # whether we can re-use "toret" workspace
-            if out is None:
-                real1 = toret
-            else:
-                real1 = RealField(pm)
-                real1[:] = 0
+            real1 = RealField(pm)
+            real1[:] = 0
 
             # the second, shifted mesh (always needed)
             real2 = RealField(pm)
@@ -407,21 +404,28 @@ def to_real_field(self, out=None, normalize=True):
                 # 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)
-                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)
+        # 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)
+            # 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"
-                if real1 is not toret:
-                    toret[:] += real1[:]
+            # need to add to the returned mesh if user supplied "out"
+            toret[:] += real1[:]
 
         # unweighted number of objects
         N = pm.comm.allreduce(Nlocal)
@@ -434,8 +438,10 @@ def to_real_field(self, out=None, normalize=True):
 
         # make sure we painted something!
         if N == 0:
-            raise ValueError(("trying to paint particle source to mesh, "
-                              "but no particles were found!"))
+            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

nbodykit/source/catalogmesh/tests/test_species.py

@@ -77,3 +77,42 @@ def test_paint(comm):
 
     # must be the same
     assert_allclose(combined.value, (real1.value + real2.value)/norm, atol=1e-5)
+
+@MPITest([1, 4])
+def test_paint_interlaced(comm):
+
+    CurrentMPIComm.set(comm)
+
+    # the test case fails only if there is enough particles to trigger
+    # the second loop of the interlaced painter; these parameters will do it.
+
+    # the catalog
+    source1 = UniformCatalog(nbar=1e-0, BoxSize=111, seed=111)
+    source2 = UniformCatalog(nbar=1e-0, BoxSize=111, seed=111)
+    source1['Weight'] = 1.0
+    source2['Weight'] = 0.1
+    cat = MultipleSpeciesCatalog(['data', 'randoms'], source1, source2, use_cache=False)
+
+    # the meshes
+    mesh = cat.to_mesh(Nmesh=32, interlaced=True)
+    mesh1 = source1.to_mesh(Nmesh=32, interlaced=True)
+    mesh2 = source2.to_mesh(Nmesh=32, interlaced=True)
+
+    # paint
+    real1 = mesh1.to_real_field()
+    real2 = mesh2.to_real_field()
+    assert_allclose(real1.cmean(), 1.0)
+    assert_allclose(real2.cmean(), 1.0)
+
+    # un-normalize real1 and real2
+    real1[:] *= real1.attrs['num_per_cell']
+    real2[:] *= real2.attrs['num_per_cell']
+    norm = real1.attrs['num_per_cell'] + real2.attrs['num_per_cell']
+
+    # the combined density field
+    #combined = mesh.to_real_field()
+    combined = mesh.paint()
+
+    assert_allclose(combined.cmean(), 1.0)
+    # must be the same
+    assert_allclose(combined.value, (real1.value + real2.value)/norm, atol=1e-5)