def box_coordinates(self, origin, diameter):
"""
Return the 8 coordinates of the box associated with the key.
Let the unit cube be described as follows:
[
[0, 0, 0],
[1, 0, 0],
[0, 1, 0],
[1, 1, 0],
[0, 0, 1],
[1, 0, 1],
[0, 1, 1],
[1, 1, 1]
]
The function returns the box coordinates in the same order
as the above unit cube coordinates.
"""
coords = np.empty(24, dtype=np.float64)
coords_data = ffi.from_buffer("double(*)[24]", coords)
origin = np.array(origin, dtype=np.float64)
origin_data = ffi.from_buffer("double(*)[3]", origin)
diameter = np.array(diameter, dtype=np.float64)
diameter_data = ffi.from_buffer("double(*)[3]", diameter)
lib.morton_key_box_coordinates(self.ctype, origin_data, diameter_data,
coords_data)
return coords.reshape(8, 3)
def from_point(cls, point, origin, diameter):
"""Create a Morton key from a point at the deepest level."""
point = np.array(point, dtype=np.float64)
point_data = ffi.from_buffer("double(*)[3]", point)
origin = np.array(origin, dtype=np.float64)
origin_data = ffi.from_buffer("double(*)[3]", origin)
diameter = np.array(diameter, dtype=np.float64)
diameter_data = ffi.from_buffer("double(*)[3]", diameter)
return cls(
lib.morton_key_from_point(point_data, origin_data, diameter_data))
def to_coordinates(self, origin, diameter):
"""Return the coordinates of the anchor."""
coords = np.empty(3, dtype=np.float64)
coords_data = ffi.from_buffer("double(*)[3]", coords)
origin = np.array(origin, dtype=np.float64)
origin_data = ffi.from_buffer("double(*)[3]", origin)
diameter = np.array(diameter, dtype=np.float64)
diameter_data = ffi.from_buffer("double(*)[3]", diameter)
lib.morton_key_to_coordinates(self.ctype, origin_data, diameter_data,
coords_data)
return coords
def from_global_points(cls, points, balanced, comm):
"""
Construct a distributed tree from a set of globally distributed points.
Params:
------
points : np.array(shape=(n_points, 3), dtype=np.float64)
Cartesian points at this processor.
balanced : bool
If 'True' constructs a balanced tree, if 'False' constructs an unbalanced tree.
comm : Intracomm
An mpi4py Intracommunicator.
"""
points = np.array(points, dtype=np.float64, order="C", copy=False)
npoints, _ = points.shape
points_data = ffi.from_buffer(f"double(*)[3]", points)
balanced_data = ffi.cast("bool", np.bool(balanced))
npoints_data = ffi.cast("size_t", npoints)
p_comm = MPI._addressof(comm)
raw_comm = ffi.cast("uintptr_t*", p_comm)
return cls(
lib.distributed_tree_from_points(points_data, npoints_data,
balanced_data, raw_comm),
comm,
p_comm,
raw_comm,
)
def find_key_in_direction(self, direction):
"""
Find a key in a given direction.
Given an integer list `direction` containing
3 elements, return the key obtained by moving
from the current key `direction[j]` steps along
dimension [j]. For example, if `direction = [2, -1, 1]`
the method returns the key by moving two boxes in positive
x-direction, one box in the negative y direction and one box
in the positive z direction. Boxes are counted with respect to
the current level.
If there is no box in the given direction, i.e. the new coordinates
are out of bounds, the method retunrs None.
"""
direction = np.array(direction, dtype=np.int64)
direction_data = ffi.from_buffer("int64_t(*)[3]", direction)
ptr = lib.morton_key_key_in_direction(self.ctype, direction_data)
if ptr == ffi.NULL:
return None
else:
return MortonKey(ptr)
def children(self):
"""Return the children."""
ptr = np.empty(8, dtype=np.uint64)
ptr_data = ffi.from_buffer("uintptr_t *", ptr)
lib.morton_key_children(self.ctype, ptr_data)
children = [
MortonKey(ffi.cast("MortonKey *", ptr[index]))
for index in range(8)
]
return children
def _clone(self, start, stop):
"""Clone a slice of the exposed Vec<T> into a Python datatype."""
n = ffi.cast("size_t", len(self))
nslice = stop - start
start = ffi.cast("size_t", start)
stop = ffi.cast("size_t", stop)
ptr = np.empty(nslice, dtype=np.uint64)
ptr_data = ffi.from_buffer("uintptr_t *", ptr)
self._iterator_protocol.clone(self._head, ptr_data, n, start, stop)
return [
self._iterator_protocol.p_type(
ffi.cast(f"{self._iterator_protocol.c_name} *", ptr[index])
)
for index in range(nslice)
]
def from_global_points(cls, points, comm):
points = np.array(points, dtype=np.float64, order="C")
npoints, _ = points.shape
points_data = ffi.from_buffer(f"double(*)[3]", points)
n_points_data = ffi.cast("size_t", npoints)
return cls(lib.domain_from_global_points(points_data, npoints, comm))
def from_anchor(cls, anchor):
"""Create a Morton key from a given anchor."""
anchor = np.array(anchor, dtype=np.uint64)
data = ffi.from_buffer("uint64_t(*)[3]", anchor)
return cls(lib.morton_key_from_anchor(data))