Add a function to compute position from ID.

nbodykit/tests/test_transform.py

@@ -188,6 +188,29 @@ def test_constarray(comm):
     a = ConstantArray([1.0, 1.0], 3, chunks=1000)
     assert a.shape == (3, 2)
 
+@MPITest([1])
+def test_idtoposition(comm):
+    ID = [0, 1, 2, 3, 4, 5, 6, 7, 8]
+    strides = [1, 4, 2]
+    sizes = 2
+    shift = 0.1
+    scale = 10.
+    pos = transform.IDToPosition(ID, strides=strides,
+                                          sizes=sizes,
+                                          shift=shift,
+                                          scale=scale)
+    numpy.testing.assert_array_equal(pos, [
+        [0.1, 0.1, 0.1],
+        [10.1, 0.1, 0.1],
+        [0.1, 0.1, 10.1],
+        [10.1, 0.1, 10.1],
+        [0.1, 10.1, 0.1],
+        [10.1, 10.1, 0.1],
+        [0.1, 10.1, 10.1],
+        [10.1, 10.1, 10.1],
+        [0.1, 0.1, 0.1],
+    ])
+
 @MPITest([1, 4])
 def test_vector_projection(comm):
     cosmo = cosmology.Planck15

nbodykit/transform.py

@@ -486,6 +486,51 @@ def mass_to_radius(mass, redshift):
 
     return da.map_blocks(mass_to_radius, mass, redshift, dtype=mass.dtype)
 
+def IDToPosition(ID, strides, scale, shift, sizes):
+    """
+    Convert ID to position assuming a regular grid labelling scheme.
+
+    ID = i * strides[0] + j * strides[1] + k * strides[2]
+    position = transpose([i, j, k]) * scale + shift
+
+    i, j, k are between `[0, sizes[.])`.
+
+    Parameters
+    ----------
+    ID : array_like
+      particle ID, starting from 0.
+    strides : array_like broadcast to shape = (3,)
+      the strides; `cat.attrs['q.strides']`
+    scale : array_like broadcast to shape = (3,)
+      the scale; `cat.attrs['q.scale']`
+    shift : array_like, broadcast to shape = (3,)
+      the shift. In the position units. `cat.attrs['q.shift']`
+    sizes: array_like broadcast to shape = (3,)
+      the size; `cat.attrs['q.sizes']` or Nc.
+
+    Returns
+    -------
+    position : position according to the formula.
+
+    """
+
+    strides, scale, shift, sizes, _ = numpy.broadcast_arrays(strides,
+                                                             scale,
+                                                             shift,
+                                                             sizes,
+                                                             numpy.ones(3))
+    def id_to_position(ID):
+        j = ID
+        qq = numpy.empty_like(ID, dtype='f8')
+        for d in range(ID.shape[-1]):
+            qq[:, d] = (ID[:, d] // strides[d]) % sizes[d]
+        return qq*scale + shift
+
+    if not isinstance(ID, da.Array):
+        ID = da.from_array(ID, chunks=100000)
+
+    ID, _  = da.broadcast_arrays(ID[:, None], numpy.ones(3))
+    return da.map_blocks(id_to_position, ID, dtype='f8')
 
 def VectorProjection(vector, direction):
     r"""
@@ -515,7 +560,6 @@ def VectorProjection(vector, direction):
 
     return projection
 
-
 # deprecated functions
 vstack = deprecate("nbodykit.transform.vstack", StackColumns, "nbodykit.transform.StackColumns")
 concatenate = deprecate("nbodykit.transform.concatenate", ConcatenateSources, "nbodykit.transform.ConcatenateSources")