import numpy
[docs]class MeshSlab(object):
"""
A convenience class to represent a specific slab of a mesh,
which is denoted as a ``slab``
"""
[docs] def __init__(self, islab, coords, axis, symmetry_axis):
"""
Parameters
----------
islab : int, [0, x[0].shape[0]]
the index of the slab, which indexes the first dimension of
the mesh (the `x` coordinate), thus producing a y-z plane
coords : list of arrays
the coordinate arrays of the mesh, with proper 3D shapes
for easy broadcasting; if the mesh has size (Nx, Ny, Nz),
then the shapes of `x` are: [(Nx, 1, 1), (1, Ny, 1), (1, 1, Nz)]
axis : int, {0, 1, 2}
the index of the mesh axis to iterate over
symmetry_axis : int, optional
if provided, the axis that has been compressed due to Hermitian symmetry
"""
self.ndim = len(coords)
self._index = islab
self.axis = axis
self.symmetry_axis = symmetry_axis
self._coords = coords
# make sure symmetry_axis > 0 (sanity check)
if self.hermitian_symmetric and self.symmetry_axis < 0:
raise ValueError("`symmetry_axis` in MeshSlab must be non-negative")
def __str__(self):
name = self.__class__.__name__
return "<%s: axis=%d, index=%d>" %(name, self.axis, self._index)
def __repr__(self):
return self.__str__()
@property
def index(self):
"""
Return an indexing list appropriate for selecting the slicing the
appropriate slab out of a full 3D array of shape (Nx, Ny, Nz)
"""
toret = [slice(None)]*self.ndim
toret[self.axis] = self._index
return tuple(toret)
@property
def meshshape(self):
"""
Return the local shape of the mesh on this rank, as determined
by the input coordinates array
"""
return tuple([numpy.shape(self._coords[i])[i] for i in range(self.ndim)])
@property
def shape(self):
"""
Return the shape of the slab
"""
return tuple([s for i, s in enumerate(self.meshshape) if i != self.axis])
@property
def hermitian_symmetric(self):
"""
Whether the slab is Hermitian-symmetric
"""
return self.symmetry_axis is not None
[docs] def coords(self, i):
"""
Return the coordinate array for dimension ``i`` on this slab,
.. note::
The return value will be properly squeezed for easy
broadcasting, i.e., if `i` is `self.axis`, then an array of
shape `(1,1)` array is returned, otherwise, the shape is
`(N_i, 1)` or `(1, N_i)`
Parameters
----------
i : int, {0,1,2}
the index of the desired dimension
Returns
-------
array_like
the coordinate array for dimension `i` on the slab; see the
note about the shape of the return array for details
"""
if i < 0: i += self.ndim
assert 0 <= i < self.ndim, "i should be between 0 and %d" %self.ndim
if i != self.axis:
return numpy.take(self._coords[i], 0, axis=self.axis)
else:
return numpy.take(self._coords[i], self._index, axis=self.axis)
[docs] def norm2(self):
"""
The square of coordinate grid norm defined at each
point on the slab.
This broadcasts the coordinate arrays along each dimension
to compute the norm at each point in the slab.
Returns
-------
array_like, (slab.shape)
the square of coordinate mesh at each point in the slab
"""
return sum(self.coords(i)**2 for i in range(self.ndim))
[docs] def mu(self, los):
"""
The `mu` value defined at each point on the slab for the
specified line-of-sight index
Parameters
----------
los: array_like,
the direction of the line-of-sight, which `mu` is defined
with respect to; must have a norm of 1.
Returns
-------
array_like, (slab.shape)
the `mu` value at each point in the slab
"""
norm = self.norm2()**0.5
with numpy.errstate(invalid='ignore', divide='ignore'):
result = sum(self.coords(i) * los[i] for i in range(self.ndim)) / norm
result[norm == 0.0] = 0.0
return result
@property
def nonsingular(self):
"""
The indices on the slab of the positive frequencies
along the dimension specified by `symmetry_axis`.
This takes advantage of the fact that :class:`pmesh.pm.ComplexField`
shifts the Nyquist frequencies to the negative halves.
Therefore this ensures that the zero and Nyquist planes perpendicular
to the symmetry axis have weight 1, whereas other modes have weight 2.
Returns
-------
idx : array_like, self.shape
Return a boolean array with the shape of the slab,
with `True` elements giving the elements with
positive freq along `symmetry_axis`
"""
try:
return self._nonsingular
except AttributeError:
# initially, return slice that includes all elements
idx = numpy.ones(self.shape, dtype=bool)
# iteration axis is symmetry axis
if self.symmetry_axis == self.axis:
# check if current iteration value is positive
if numpy.float(self.coords(self.axis)) <= 0.:
idx = numpy.zeros(self.shape, dtype=bool)
# one of slab dimensions is symmetry axis
else:
# get the indices that have positive freq along symmetry axis
nonsingular = (self._coords[self.symmetry_axis] > 0.)
nonsingular = numpy.take(nonsingular, 0, axis=self.axis)
idx[...] = nonsingular
self._nonsingular = idx
return self._nonsingular
@property
def hermitian_weights(self):
"""
Weights to be applied to quantities on the slab in order
to account for Hermitian symmetry
These weights double-count the positive frequencies along
the `symmetry_axis`.
"""
try:
return self._weights
except AttributeError:
# if not Hermitian symmetric, weights are 1
if not self.hermitian_symmetric:
toret = 1.
# iteration axis is symmetry axis
elif self.axis == self.symmetry_axis:
toret = 1.
if numpy.float(self.coords(self.symmetry_axis)) > 0.:
toret = 2.
# only nonsingular plane gets factor of 2
else:
toret = numpy.ones(self.shape, dtype='f4')
toret[self.nonsingular] = 2.
self._weights = toret
return self._weights
[docs]def SlabIterator(coords, axis=0, symmetry_axis=None):
"""
Iterate over the specified dimension of the coordinate mesh,
returning a :class:`MeshSlab` for each iteration
Parameters
----------
coords : list of arrays
the coordinate arrays of the mesh, with proper 3D shapes
for easy broadcasting; if the mesh has size (Nx, Ny, Nz),
then the shapes of `x3d` should be: ``[(Nx, 1, 1), (1, Ny, 1), (1, 1, Nz)]``
axis : int, optional
the index of the mesh axis to iterate over
symmetry_axis : int, optional
if provided, the axis that has been compressed due to Hermitian symmetry
"""
# number of dimensions in the mesh
ndim = len(coords)
if ndim not in [2,3]:
raise NotImplementedError("SlabIterator can only be used on 3D or 2D arrays")
# account for negative axes
if axis < 0: axis += ndim
assert 0 <= axis < ndim, "axis should be between 0 and %d" %ndim
if symmetry_axis is not None and symmetry_axis < 0: symmetry_axis += ndim
# this will only work if shapes are: [(Nx, 1, 1), (1, Ny, 1), (1, 1, Nz)]
# mainly a sanity check to make sure things work
shapes = [numpy.shape(x) for x in coords]
try:
mesh_size = [shape[i] for i, shape in enumerate(shapes)]
for i in range(ndim):
desired_shape = numpy.ones(ndim)
desired_shape[i] = mesh_size[i]
if shapes[i] != tuple(desired_shape):
raise ValueError("coordinate array shape mismatch")
except:
raise ValueError("input coordinates with shapes %s are not correct" %str(shapes))
# iterate over the specified axis, slab by slab
N = numpy.shape(coords[axis])[axis]
for islab in range(N):
yield MeshSlab(islab, coords, axis, symmetry_axis)