diff --git a/nbodykit/tests/test_transform.py b/nbodykit/tests/test_transform.py index aa5695c2..6ec74fb5 100644 --- a/nbodykit/tests/test_transform.py +++ b/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 diff --git a/nbodykit/transform.py b/nbodykit/transform.py index a98c1edd..f831e279 100644 --- a/nbodykit/transform.py +++ b/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")