WIP: Improving the Grid to support reading k and Fourier resampling

examples/paint/test_fastpm.params

@@ -1,10 +1,10 @@
-Nmesh: 64
 output: ${NBKIT_HOME}/examples/output/test_paint_fastpm
-
+Nmesh : 64
+paintNmesh: 128
 DataSource:
    plugin: FastPM
    path: ${NBKIT_CACHE}/data/fastpm_1.0000
-
+Nfile : 4
 Painter:
    plugin: DefaultPainter
 

examples/paint/test_grid.params

@@ -1,10 +1,16 @@
 Nmesh: 64
+paintNmesh : 128
 output: ${NBKIT_HOME}/examples/output/test_paint_grid
 nfile : 4
 DataSource:
    plugin: BigFileGrid
    path: ${NBKIT_CACHE}/data/bigfile_grid
    dataset: PaintGrid
+
 Painter:
    plugin: DefaultPainter
+   normalize : true
+   setMean : 1.0   # useful for over density inputs.
+   fk : exp(- 0.5 * (k * 8) ** 2)
+   frho : (rho - 1) ** 2
 

examples/paint/test_grid_downsample.params

@@ -0,0 +1,10 @@
+Nmesh: 32
+output: ${NBKIT_HOME}/examples/output/test_paint_grid_32
+nfile : 4
+DataSource:
+   plugin: BigFileGrid
+   path: ${NBKIT_CACHE}/data/bigfile_grid
+   dataset: PaintGrid
+Painter:
+   plugin: DefaultPainter
+

examples/paint/test_grid_upsample.params

@@ -0,0 +1,10 @@
+Nmesh: 128
+output: ${NBKIT_HOME}/examples/output/test_paint_grid_128
+nfile : 4
+DataSource:
+   plugin: BigFileGrid
+   path: ${NBKIT_CACHE}/data/bigfile_grid
+   dataset: PaintGrid
+Painter:
+   plugin: DefaultPainter
+

examples/paint/test_gridk.params

@@ -0,0 +1,11 @@
+Nmesh: 64
+paintNmesh : 128
+output: ${NBKIT_HOME}/examples/output/test_paint_gridk
+nfile : 4
+DataSource:
+   plugin: BigFileGrid
+   path: ${NBKIT_CACHE}/data/bigfile_gridk
+   dataset: LinearDensityK
+Painter:
+   plugin: DefaultPainter
+

nbodykit/plugins/algorithms/PaintGrid.py

@@ -1,5 +1,7 @@
 from nbodykit.extensionpoints import Algorithm
 from nbodykit.extensionpoints import DataSource, GridSource, Painter
+from nbodykit import resampler
+
 import os
 import numpy
 
@@ -14,13 +16,20 @@ class PaintGridAlgorithm(Algorithm):
     """
     plugin_name = "PaintGrid"
 
-    def __init__(self, Nmesh, DataSource, Painter=None, paintbrush='cic', dataset='PaintGrid', Nfile=0):
+    def __init__(self, Nmesh, DataSource, Painter=None, paintbrush='cic', paintNmesh=None, dataset='PaintGrid', Nfile=0, writeFourier=False):
         # combine the two fields
         self.datasource = DataSource
         if Painter is None:
             Painter = Painter.create("DefaultPainter")
         self.painter = Painter
         self.dataset = dataset
+        if paintNmesh is None:
+            self.paintNmesh = self.Nmesh
+
+        if self.Nfile == 0:
+            chunksize = 1024 * 1024 * 512
+            self.Nfile = (self.Nmesh * self.Nmesh * self.Nmesh + chunksize - 1)// chunksize
+
 
     @classmethod
     def register(cls):
@@ -39,7 +48,9 @@ def register(cls):
 
         s.add_argument('dataset', help="name of dataset to write to")
         s.add_argument('Nfile', required=False, help="number of files")
-
+        s.add_argument('writeFourier', type=bool, required=False, help="Write complex Fourier modes instead?")
+        s.add_argument('paintNmesh', type=int, required=False,
+                    help="The painting Nmesh. The grid will be Fourier resampled to Nmesh before output. A value larger than Nmesh can reduce grid artifacts.")
         s.add_argument('paintbrush', type=lambda x: x.lower(), choices=['cic', 'tsc'],
             help='the density assignment kernel to use when painting; '
                  'CIC (2nd order) or TSC (3rd order)')
@@ -53,9 +64,10 @@ def run(self):
 
         if self.comm.rank == 0:
             self.logger.info('importing done')
-
+            self.logger.info('Resolution Nmesh : %d' % self.paintNmesh)
+            self.logger.info('paintbrush : %s' % self.paintbrush)
         # setup the particle mesh object, taking BoxSize from the painters
-        pm = ParticleMesh(self.datasource.BoxSize, self.Nmesh, 
+        pm = ParticleMesh(self.datasource.BoxSize, self.paintNmesh,
                             paintbrush=self.paintbrush, dtype='f4', comm=self.comm)
 
         stats = self.painter.paint(pm, self.datasource)
@@ -77,32 +89,21 @@ def save(self, output, result):
         """
         import bigfile
         import numpy
-        import mpsort
         pm, stats = result
-        x3d = pm.real.copy()
-        istart = pm.partition.local_i_start
-        ind = numpy.zeros(x3d.shape, dtype='i8')
-        for d in range(3):
-            i = numpy.arange(istart[d], istart[d] + x3d.shape[d])
-            i = i.reshape([-1 if dd == d else 1 for dd in range(3)])
-            ind[...] *= pm.Nmesh
-            ind[...] += i
-
-        x3d = x3d.ravel()
-        ind = ind.ravel()
-        mpsort.sort(x3d, orderby=ind, comm=self.comm)
-        if self.Nfile == 0:
-            chunksize = 1024 * 1024 * 512
-            Nfile = (self.Nmesh * self.Nmesh * self.Nmesh + chunksize - 1)// chunksize
-        else:
-            Nfile = self.Nfile
 
         if self.comm.rank == 0:
-            self.logger.info("writing to %s/%s in %d parts" % (output, self.dataset, Nfile))
+            self.logger.info('Output Nmesh : %d' % self.Nmesh)
+            self.logger.info("writing to %s/%s in %d parts" % (output, self.dataset, self.Nfile))
 
         f = bigfile.BigFileMPI(self.comm, output, create=True)
-        b = f.create_from_array(self.dataset, x3d, Nfile=Nfile)
+
+        array = resampler.write(pm, self.Nmesh, self.writeFourier)
+
+        b = f.create_from_array(self.dataset, array, Nfile=self.Nfile)
+
         b.attrs['ndarray.shape'] = numpy.array([self.Nmesh, self.Nmesh, self.Nmesh], dtype='i8')
         b.attrs['BoxSize'] = numpy.array(self.datasource.BoxSize, dtype='f8')
         b.attrs['Nmesh'] = self.Nmesh
+        b.attrs['paintNmesh'] = self.paintNmesh
+        b.attrs['paintbrush'] = self.paintbrush
         b.attrs['Ntot'] = stats['Ntot']

nbodykit/plugins/datasource/BigFileGrid.py

@@ -1,68 +1,50 @@
 from nbodykit.extensionpoints import GridSource
+from nbodykit import resampler
 import numpy
 
 class BigFileGridSource(GridSource):
     """
     Class to read field gridded data from a binary file
-    
+
     Notes
     -----
     * Reading is designed to be done by `GridPainter`, which
       reads gridded quantity straight into the `ParticleMesh`
     """
     plugin_name = "BigFileGrid"
-    
+
     def __init__(self, path, dataset):
         import bigfile
         f = bigfile.BigFileMPI(self.comm, self.path)
-        self.dataset = dataset
-        self.path = path
+
         with f[self.dataset] as d:
             self.BoxSize = d.attrs['BoxSize']
-            self.Nmesh = d.attrs['Nmesh'][0]
-            self.Ntot = d.attrs['Ntot'][0]
+            self.Nmesh = int(d.attrs['Nmesh'][0])
+            if 'Ntot' in d.attrs:
+                self.Ntot = d.attrs['Ntot'][0]
+            else:
+                self.Ntot = 0
+
+            # Is this a complex field or a real field?
+            if d.dtype.kind == 'c':
+                self.isfourier = True
+            else:
+                self.isfourier = False
 
     @classmethod
     def register(cls):
         s = cls.schema
-        s.description = "read gridded field data from a binary file"
+        s.description = "read gridded field data from a binary file. Fourier resampling is applied if necessary."
 
         s.add_argument("path", type=str, help="the file path to load the data from")
         s.add_argument("dataset", type=str, help="the file path to load the data from")
 
     def read(self, pm):
-        if pm.Nmesh != self.Nmesh:
-            raise ValueError("The Grid has a mesh size of %d, but ParticleMesh has %d" % (self.Nmesh, pm.Nmesh))
-
         import bigfile
-        import mpsort
-        f = bigfile.BigFileMPI(self.comm, self.path)
-
-        istart = pm.partition.local_i_start
-
-        x3d = numpy.empty(pm.real.shape, dtype='f4')
-
-        ind = numpy.zeros(x3d.shape, dtype='i8')
-
-        for d in range(3):
-            i = numpy.arange(istart[d], istart[d] + x3d.shape[d])
-            i = i.reshape([-1 if dd == d else 1 for dd in range(3)])
-            ind[...] *= pm.Nmesh
-            ind[...] += i
-
-        ind = ind.ravel()
-        x3d = x3d.ravel()
-
-        start = sum(self.comm.allgather(x3d.size)[:self.comm.rank])
-        end = start + x3d.size
-
-        originind = numpy.arange(start, end, dtype='i8')
-
-        mpsort.sort(originind, orderby=ind, comm=self.comm)
+        if self.comm.rank == 0:
+            self.logger.info("Reading from Nmesh = %d to Nmesh = %d" %(self.Nmesh, pm.Nmesh))
 
+        f = bigfile.BigFileMPI(self.comm, self.path)
         with f[self.dataset] as ds:
-            x3d[:] = ds[start: end]
-
-        mpsort.sort(x3d, orderby=originind, comm=self.comm)
+            resampler.read(pm, ds, self.Nmesh, self.isfourier)
 
-        pm.real.flat[:] = x3d

nbodykit/plugins/painter/DefaultPainter.py

@@ -4,17 +4,22 @@
 class DefaultPainter(Painter):
     plugin_name = "DefaultPainter"
 
-    def __init__(self, weight=None):
+    def __init__(self, weight=None, frho=None, fk=None, normalize=False, setMean=None):
         pass
 
     @classmethod
     def register(cls):
         s = cls.schema
         s.description = "grid the density field of an input DataSource of objects, optionally "
-        s.description += "using a weight for each object"
+        s.description += "using a weight for each object. "
 
         s.add_argument("weight", help="the column giving the weight for each object")
 
+        s.add_argument("frho", type=str, help="A python expresion for transforming the real space density field. variables: rho. example: 1 + (rho - 1)**2")
+        s.add_argument("fk", type=str, help="A python expresion for transforming the fourier space density field. variables: k, kx, ky, kz. example: exp(-(k * 0.5)**2). applied before frho ")
+        s.add_argument("normalize", type=bool, help="Normalize the field to set mean == 1. Applied before fk.")
+        s.add_argument("setMean", type=float, help="Set the mean. Applied after normalize.")
+
     def paint(self, pm, datasource):
         """
         Paint the ``DataSource`` specified by ``input`` onto the 
@@ -52,4 +57,44 @@ def paint(self, pm, datasource):
         elif isinstance(datasource, GridSource):
             datasource.read(pm)
             stats['Ntot'] = datasource.Ntot
+
+        # apply the filters.
+
+        mean = self.comm.allreduce(pm.real.sum(dtype='f8')) / pm.Nmesh ** 3.
+
+        if self.comm.rank == 0:
+            self.logger.info("Mean = %g" % mean)
+
+        if self.normalize:
+            pm.real *= 1. / mean
+            mean = self.comm.allreduce(pm.real.sum(dtype='f8')) / pm.Nmesh ** 3.
+            if self.comm.rank == 0:
+                self.logger.info("Renormalized mean = %g" % mean)
+
+        if self.setMean is not None:
+            pm.real += (self.setMean - mean)
+
+        if self.fk:
+            if self.comm.rank == 0:
+                self.logger.info("applying transformation fk %s" % self.fk)
+
+            def function(k, kx, ky, kz):
+                from numpy import exp, sin, cos
+                return eval(self.fk)
+            pm.r2c()
+            k = (pm.k[0] ** 2 + pm.k[1] ** 2 + pm.k[2] ** 2) ** 0.5
+            pm.complex[...] *= function(k, pm.k[0], pm.k[1], pm.k[2])
+            pm.c2r()
+
+        if self.frho:
+            if self.comm.rank == 0:
+                self.logger.info("applying transformation frho %s" % self.frho)
+
+            def function(rho):
+                return eval(self.frho)
+            if self.comm.rank == 0:
+                self.logger.info("example value before frho %g" % pm.real.flat[0])
+            pm.real[...] = function(pm.real)
+            if self.comm.rank == 0:
+                self.logger.info("example value after frho %g" % pm.real.flat[0])
         return stats