Source code for nbodykit.source.mesh.array

from nbodykit.base.mesh import MeshSource

from nbodykit import CurrentMPIComm
from nbodykit.utils import attrs_to_dict
from pmesh.pm import RealField, ComplexField
import numpy

[docs]class ArrayMesh(MeshSource): """ A MeshSource initalized from an in-memory numpy array. .. note:: The in-memory array must be fully hosted by the root rank. Parameters ---------- array : numpy.ndarray the numpy array holding the field data; this must be fully hosted by the rank specified by ``root`` BoxSize : float, 3-vector the size of the box root : int, optional the root rank holding the array data **kwargs : additional meta-data to store """ def __repr__(self): return "ArrayMesh()" @CurrentMPIComm.enable def __init__(self, array, BoxSize, comm=None, root=0, **kwargs): if comm.rank == root: array = numpy.array(array) if array.dtype.kind == 'c': # transform to real for the correct shape array = numpy.fft.irfftn(array) array[...] *= numpy.prod(array.shape) shape = array.shape dtype = array.dtype else: array, dtype, shape = [None] * 3 dtype = comm.bcast(dtype, root=root) shape = comm.bcast(shape, root=root) assert len(shape) in (2, 3) Nmesh = shape empty = numpy.empty((0,), dtype) MeshSource.__init__(self, comm, Nmesh, BoxSize, empty.real.dtype) self.field = self.pm.create(type='real') if comm.rank != root: array = empty # ignore data from other ranks. else: array = array.ravel() # fill the field with the array self.field.unravel(array)
[docs] def to_real_field(self): if isinstance(self.field, RealField): return self.field.copy() else: return NotImplemented