allowing for input scattered array in ArrayMesh

nb/window_examples.ipynb

@@ -219,7 +219,7 @@
     "boxsize = 10000.\n",
     "edges = {'step': 2.*np.pi/boxsize}\n",
     "window_large = CatalogSmoothWindow(randoms_positions1=randoms_positions, randoms_weights1=randoms_weights,\n",
-    "                              power_ref=poles, edges=edges, boxsize=boxsize, position_type='pos', dtype='f4').poles\n",
+    "                                   power_ref=poles, edges=edges, boxsize=boxsize, position_type='pos', dtype='f4').poles\n",
     "# You can save the window function\n",
     "with tempfile.TemporaryDirectory() as tmp_dir:\n",
     "    fn = os.path.join(tmp_dir, 'tmp.npy')\n",
@@ -319,7 +319,7 @@
    ],
    "source": [
     "window_small = CatalogSmoothWindow(randoms_positions1=randoms_positions, randoms_weights1=randoms_weights,\n",
-    "                              power_ref=poles, edges=edges, boxsize=3000., position_type='pos', dtype='f4').poles\n",
+    "                                   power_ref=poles, edges=edges, boxsize=3000., position_type='pos', dtype='f4').poles\n",
     "\n",
     "window = PowerSpectrumSmoothWindow.concatenate_x(window_large, window_small, frac_nyq=0.9) # Here we remove the range above 0.9 Nyquist (which may not be reliable)\n",
     "\n",

pypower/fft_power.py

@@ -1938,16 +1938,16 @@ def __init__(self, data_positions1, data_positions2=None, randoms_positions1=Non
         nmesh : array, int, default=None
             Mesh size, i.e. number of mesh nodes along each axis.
 
-        boxsize : float, default=None
-            Physical size of the box, defaults to maximum extent taken by all input positions, times ``boxpad``.
+        boxsize : array, float, default=None
+            Physical size of the box along each axis, defaults to maximum extent taken by all input positions, times ``boxpad``.
 
         boxcenter : array, float, default=None
             Box center, defaults to center of the Cartesian box enclosing all input positions.
 
         cellsize : array, float, default=None
             Physical size of mesh cells.
             If not ``None``, and mesh size ``nmesh`` is not ``None``, used to set ``boxsize`` as ``nmesh * cellsize``.
-            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize/cellsize``.
+            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize / cellsize``.
 
         boxpad : float, default=2.
             When ``boxsize`` is determined from input positions, take ``boxpad`` times the smallest box enclosing positions as ``boxsize``.

pypower/fft_window.py

@@ -820,8 +820,8 @@ def __init__(self, randoms_positions1=None, randoms_positions2=None,
             Mesh size, i.e. number of mesh nodes along each axis.
             If ``None``, defaults to the value used in estimation of ``power_ref``.
 
-        boxsize : float, default=None
-            Physical size of the box, defaults to maximum extent taken by all input positions, times ``boxpad``.
+        boxsize : array, float, default=None
+            Physical size of the box along each axis.
             If ``None``, defaults to the value used in estimation of ``power_ref``.
 
         boxcenter : array, float, default=None
@@ -831,7 +831,7 @@ def __init__(self, randoms_positions1=None, randoms_positions2=None,
         cellsize : array, float, default=None
             Physical size of mesh cells.
             If not ``None``, and mesh size ``nmesh`` is not ``None``, used to set ``boxsize`` as ``nmesh * cellsize``.
-            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize/cellsize``.
+            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize / cellsize``.
 
         boxpad : float, default=2.
             When ``boxsize`` is determined from input positions, take ``boxpad`` times the smallest box enclosing positions as ``boxsize``.

pypower/mesh.py

@@ -130,8 +130,8 @@ def _get_mesh_attrs(nmesh=None, boxsize=None, boxcenter=None, cellsize=None, pos
     cellsize : array, float, default=None
         Physical size of mesh cells.
         If not ``None``, ``boxsize`` is ``None`` and mesh size ``nmesh`` is not ``None``, used to set ``boxsize`` to ``nmesh * cellsize``.
-        If ``nmesh`` is ``None``, it is set to (the nearest integer(s) to) ``boxsize/cellsize`` if ``boxsize`` is provided,
-        else to the nearest even integer to ``boxsize/cellsize``, and ``boxsize`` is then reset to ``nmesh * cellsize``.
+        If ``nmesh`` is ``None``, it is set to (the nearest integer(s) to) ``boxsize / cellsize`` if ``boxsize`` is provided,
+        else to the nearest even integer to ``boxsize / cellsize``, and ``boxsize`` is then reset to ``nmesh * cellsize``.
 
     positions : (list of) (N, 3) arrays, default=None
         If ``boxsize`` and / or ``boxcenter`` is ``None``, use this (list of) position arrays
@@ -193,7 +193,7 @@ def _get_mesh_attrs(nmesh=None, boxsize=None, boxcenter=None, cellsize=None, pos
     return nmesh, boxsize, boxcenter
 
 
-def ArrayMesh(array, boxsize, mpiroot=0, mpicomm=MPI.COMM_WORLD):
+def ArrayMesh(array, boxsize, nmesh=None, mpiroot=0, mpicomm=MPI.COMM_WORLD):
     """
     Turn numpy array into :class:`pmesh.pm.RealField`.
 
@@ -202,11 +202,15 @@ def ArrayMesh(array, boxsize, mpiroot=0, mpicomm=MPI.COMM_WORLD):
     array : array
         Mesh numpy array gathered on ``mpiroot``.
 
-    boxsize : array
-        Physical box size.
+    boxsize : array, float, default=None
+        Physical size of the box along each axis.
+
+    nmesh : array, int, default=None
+        If ``mpiroot`` is ``None``, mesh size, i.e. number of mesh nodes along each axis.
 
     mpiroot : int, default=0
         MPI rank where input array is gathered.
+        If input array is scattered accross all ranks in C ordering, pass ``mpiroot = None`` and specify ``nmesh``.
 
     mpicomm : MPI communicator, default=MPI.COMM_WORLD
         The MPI communicator.
@@ -215,20 +219,28 @@ def ArrayMesh(array, boxsize, mpiroot=0, mpicomm=MPI.COMM_WORLD):
     -------
     mesh : pmesh.pm.RealField
     """
-    if mpicomm.rank == mpiroot:
-        dtype, shape = array.dtype, array.shape
+    if mpiroot is None:
+        dtype = array.dtype
+        if nmesh is None:
+            raise ValueError('In case input mesh is scattered accross all ranks, provide its shape (nmesh)')
+        shape = _make_array(nmesh, 3, dtype='i8')
     else:
-        dtype, shape, array = None, None, None
+        if mpicomm.rank == mpiroot:
+            dtype, shape = array.dtype, array.shape
+        else:
+            dtype, shape, array = None, None, None
+        dtype = mpicomm.bcast(dtype, root=mpiroot)
+        shape = mpicomm.bcast(shape, root=mpiroot)
 
-    dtype = mpicomm.bcast(dtype, root=mpiroot)
-    shape = mpicomm.bcast(shape, root=mpiroot)
     boxsize = _make_array(boxsize, 3, dtype='f8')
     pm = ParticleMesh(BoxSize=boxsize, Nmesh=shape, dtype=dtype, comm=mpicomm)
     mesh = pm.create(type='real')
-    if mpicomm.rank == mpiroot:
-        array = array.ravel()  # ignore data from other ranks
+
+    if mpiroot is None or mpicomm.rank == mpiroot:
+        array = np.ravel(array)  # ignore data from other ranks
     else:
-        array = np.empty((0,), dtype)
+        array = np.empty((0,), dtype=dtype)
+
     mesh.unravel(array)
     return mesh
 
@@ -273,20 +285,17 @@ def __init__(self, data_positions, data_weights=None, randoms_positions=None, ra
 
         nmesh : array, int, default=None
             Mesh size, i.e. number of mesh nodes along each axis.
-            If not provided, see ``value``.
 
-        boxsize : float, default=None
-            Physical size of the box.
-            If not provided, see ``positions``.
+        boxsize : array, float, default=None
+            Physical size of the box along each axis, defaults to maximum extent taken by all input positions, times ``boxpad``.
 
         boxcenter : array, float, default=None
-            Box center.
-            If not provided, see ``positions``.
+            Box center, defaults to center of the Cartesian box enclosing all input positions.
 
         cellsize : array, float, default=None
             Physical size of mesh cells.
             If not ``None``, and mesh size ``nmesh`` is not ``None``, used to set ``boxsize`` as ``nmesh * cellsize``.
-            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize/cellsize``.
+            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize / cellsize``.
 
         wrap : bool, default=False
             Whether to wrap input positions?

pypower/smooth_window.py

@@ -576,8 +576,8 @@ def __init__(self, randoms_positions1=None, randoms_positions2=None,
             Mesh size, i.e. number of mesh nodes along each axis.
             If ``None``, defaults to the value used in estimation of ``power_ref``.
 
-        boxsize : float, default=None
-            Physical size of the box, defaults to maximum extent taken by all input positions, times ``boxpad``.
+        boxsize : array, float, default=None
+            Physical size of the box along each axis.
             If ``None``, defaults to the value used in estimation of ``power_ref``.
 
         boxcenter : array, float, default=None
@@ -587,7 +587,7 @@ def __init__(self, randoms_positions1=None, randoms_positions2=None,
         cellsize : array, float, default=None
             Physical size of mesh cells.
             If not ``None``, and mesh size ``nmesh`` is not ``None``, used to set ``boxsize`` as ``nmesh * cellsize``.
-            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize/cellsize``.
+            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize / cellsize``.
 
         boxpad : float, default=2.
             When ``boxsize`` is determined from input positions, take ``boxpad`` times the smallest box enclosing positions as ``boxsize``.

pypower/tests/test_fft_power.py

@@ -8,7 +8,7 @@
 from mockfactory import LagrangianLinearMock, Catalog
 from mockfactory.make_survey import RandomBoxCatalog
 
-from pypower import MeshFFTPower, CatalogFFTPower, CatalogMesh, PowerSpectrumStatistics, mpi, utils, setup_logging
+from pypower import MeshFFTPower, CatalogFFTPower, CatalogMesh, ArrayMesh, PowerSpectrumStatistics, mpi, utils, setup_logging
 from pypower.fft_power import normalization, normalization_from_nbar, find_unique_edges, get_real_Ylm, project_to_basis
 
 
@@ -551,6 +551,18 @@ def test_normalization():
     assert np.allclose(new, old, atol=0, rtol=1e-1)
 
 
+def test_array_mesh():
+
+    shape = (60, 50, 20)
+    array = np.arange(np.prod(shape, dtype='i'), dtype='f8')
+    array.shape = shape
+    mesh_root = ArrayMesh(array, boxsize=1., mpiroot=0)
+    mpicomm = mesh_root.pm.comm
+    array = np.ravel(array)
+    mesh_scattered = ArrayMesh(array[mpicomm.rank * array.size // mpicomm.size:(mpicomm.rank + 1) * array.size // mpicomm.size], boxsize=1., nmesh=shape, mpiroot=None)
+    assert np.allclose(mesh_scattered.value, mesh_root.value)
+
+
 def test_catalog_mesh():
 
     data = Catalog.read(data_fn)
@@ -890,6 +902,7 @@ def run(interlacing=2):
     test_power_statistic()
     test_find_edges()
     test_ylm()
+    test_array_mesh()
     test_catalog_mesh()
     test_field_power()
     test_mesh_power()

pypower/wide_angle.py

@@ -308,7 +308,7 @@ def select_proj(self, projsin=None, projsout=None, **kwargs):
                             new.append(kwargs.get(name, old[0]))
                     setattr(self, name, new)
 
-        self.value = []
+        value = []
         for iin, projin in enumerate(self.projsin):
             line = []
             for iout, projout in enumerate(self.projsout):
@@ -323,8 +323,8 @@ def select_proj(self, projsin=None, projsout=None, **kwargs):
                     else:
                         tmp = np.zeros(shape, dtype=self.dtype)
                 line.append(tmp)
-            self.value.append(line)
-        self.value = np.bmat(self.value).A
+            value.append(line)
+        self.value = np.bmat(value).A
 
     def __getitem__(self, slices):
         """Call :meth:`slice_x`."""
@@ -355,7 +355,7 @@ def slice_x(self, slicein=None, sliceout=None, projsin=None, projsout=None):
             List of output projections to apply slicing to.
             Defaults to :attr:`projsout`.
         """
-        self.value = self.unpacked()  # unpack first, as based on :attr:`xin`, :attr:`xout`
+        value = self.unpacked()  # unpack first, as based on :attr:`xin`, :attr:`xout`
 
         inprojs, masks, factors = {}, {}, {}
         for axis in ['in', 'out']:
@@ -398,9 +398,9 @@ def slice_x(self, slicein=None, sliceout=None, projsin=None, projsout=None):
             selfiin = self.projsin.index(projin)
             for iout, projout in enumerate(inprojs['out']):
                 selfiout = self.projsout.index(projout)
-                self.value[selfiin][selfiout] = self.value[selfiin][selfiout][np.ix_(masks['in'][iin], masks['out'][iout])]
+                value[selfiin][selfiout] = value[selfiin][selfiout][np.ix_(masks['in'][iin], masks['out'][iout])]
 
-        self.value = np.bmat(self.value).A
+        self.value = np.bmat(value).A
         if not all(f == 1 for f in factors.values()):
             self.rebin_x(factorin=factors['in'], factorout=factors['out'], projsin=inprojs['in'], projsout=inprojs['out'])
 
@@ -427,7 +427,7 @@ def select_x(self, xinlim=None, xoutlim=None, projsin=None, projsout=None):
             Defaults to :attr:`projsout`.
         """
         # One could also set the slices, and call slice_x, but this is inefficient in case x is different for each projection
-        self.value = self.unpacked()  # unpack first, as based on :attr:`xin`, :attr:`xout`
+        value = self.unpacked()  # unpack first, as based on :attr:`xin`, :attr:`xout`
 
         inprojs, masks = {}, {}
         for axis in ['in', 'out']:
@@ -461,9 +461,9 @@ def select_x(self, xinlim=None, xoutlim=None, projsin=None, projsout=None):
             selfiin = self.projsin.index(projin)
             for iout, projout in enumerate(inprojs['out']):
                 selfiout = self.projsout.index(projout)
-                self.value[selfiin][selfiout] = self.value[selfiin][selfiout][np.ix_(masks['in'][iin], masks['out'][iout])]
+                value[selfiin][selfiout] = value[selfiin][selfiout][np.ix_(masks['in'][iin], masks['out'][iout])]
 
-        self.value = np.bmat(self.value).A
+        self.value = np.bmat(value).A
 
     def rebin_x(self, factorin=1, factorout=1, projsin=None, projsout=None, statistic=None):
         """
@@ -495,7 +495,7 @@ def rebin_x(self, factorin=1, factorout=1, projsin=None, projsout=None, statisti
             self.rebin_x(factorin=factorin, factorout=1, projsin=projsin, projsout=projsout, statistic=np.sum)
             return
 
-        self.value = self.unpacked()  # unpack first, as based on :attr:`xin`, :attr:`xout`
+        value = self.unpacked()  # unpack first, as based on :attr:`xin`, :attr:`xout`
 
         inprojs, old_weights, new_weights = {}, {}, {}
         for axis in ['in', 'out']:
@@ -532,15 +532,15 @@ def rebin_x(self, factorin=1, factorout=1, projsin=None, projsout=None, statisti
             selfiin = self.projsin.index(projin)
             for iout, projout in enumerate(inprojs['out']):
                 selfiout = self.projsout.index(projout)
-                tmp = self.value[selfiin][selfiout]
+                tmp = value[selfiin][selfiout]
                 new_shape = tuple(s // f for s, f in zip(tmp.shape, (factorin, factorout)))
                 oweights, nweights = 1., 1.
                 for iaxis, (axis, ii) in enumerate(zip(['in', 'out'], [iin, iout])):
                     if axis in old_weights:
                         oweights = oweights * np.expand_dims(old_weights[axis][ii], axis=1 - iaxis)
                         nweights = nweights * np.expand_dims(new_weights[axis][ii], axis=1 - iaxis)
-                self.value[selfiin][selfiout] = utils.rebin(self.value[selfiin][selfiout] * oweights, new_shape, statistic=statistic) / nweights
-        self.value = np.bmat(self.value).A
+                value[selfiin][selfiout] = utils.rebin(value[selfiin][selfiout] * oweights, new_shape, statistic=statistic) / nweights
+        self.value = np.bmat(value).A
 
     @classmethod
     def concatenate_proj(cls, *others, axis='in'):
@@ -939,7 +939,7 @@ def run(self):
         """
         k = self.xin[0]
         eye = np.eye(len(k), dtype=k.dtype)
-        self.value = []
+        value = []
 
         for projin in self.projsin:
             line = []
@@ -978,7 +978,7 @@ def run(self):
                             tmp[-1, -2] = 2. * coeff / deltak[-1]
                             block += tmp.T  # (in, out)
                 line.append(block)
-            self.value.append(line)
-        self.value = np.bmat(self.value).A  # (in, out)
+            value.append(line)
+        self.value = np.bmat(value).A  # (in, out)
 
     propose_out = CorrelationFunctionOddWideAngleMatrix.propose_out