selections.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435
  1. # This file is part of h5py, a Python interface to the HDF5 library.
  2. #
  3. # http://www.h5py.org
  4. #
  5. # Copyright 2008-2013 Andrew Collette and contributors
  6. #
  7. # License: Standard 3-clause BSD; see "license.txt" for full license terms
  8. # and contributor agreement.
  9. """
  10. High-level access to HDF5 dataspace selections
  11. """
  12. import numpy as np
  13. from .base import product
  14. from .. import h5s, h5r, _selector
  15. def select(shape, args, dataset=None):
  16. """ High-level routine to generate a selection from arbitrary arguments
  17. to __getitem__. The arguments should be the following:
  18. shape
  19. Shape of the "source" dataspace.
  20. args
  21. Either a single argument or a tuple of arguments. See below for
  22. supported classes of argument.
  23. dataset
  24. A h5py.Dataset instance representing the source dataset.
  25. Argument classes:
  26. Single Selection instance
  27. Returns the argument.
  28. numpy.ndarray
  29. Must be a boolean mask. Returns a PointSelection instance.
  30. RegionReference
  31. Returns a Selection instance.
  32. Indices, slices, ellipses, MultiBlockSlices only
  33. Returns a SimpleSelection instance
  34. Indices, slices, ellipses, lists or boolean index arrays
  35. Returns a FancySelection instance.
  36. """
  37. if not isinstance(args, tuple):
  38. args = (args,)
  39. # "Special" indexing objects
  40. if len(args) == 1:
  41. arg = args[0]
  42. if isinstance(arg, Selection):
  43. if arg.shape != shape:
  44. raise TypeError("Mismatched selection shape")
  45. return arg
  46. elif isinstance(arg, np.ndarray) and arg.dtype.kind == 'b':
  47. if arg.shape != shape:
  48. raise TypeError("Boolean indexing array has incompatible shape")
  49. return PointSelection.from_mask(arg)
  50. elif isinstance(arg, h5r.RegionReference):
  51. if dataset is None:
  52. raise TypeError("Cannot apply a region reference without a dataset")
  53. sid = h5r.get_region(arg, dataset.id)
  54. if shape != sid.shape:
  55. raise TypeError("Reference shape does not match dataset shape")
  56. return Selection(shape, spaceid=sid)
  57. if dataset is not None:
  58. selector = dataset._selector
  59. else:
  60. space = h5s.create_simple(shape)
  61. selector = _selector.Selector(space)
  62. return selector.make_selection(args)
  63. class Selection(object):
  64. """
  65. Base class for HDF5 dataspace selections. Subclasses support the
  66. "selection protocol", which means they have at least the following
  67. members:
  68. __init__(shape) => Create a new selection on "shape"-tuple
  69. __getitem__(args) => Perform a selection with the range specified.
  70. What args are allowed depends on the
  71. particular subclass in use.
  72. id (read-only) => h5py.h5s.SpaceID instance
  73. shape (read-only) => The shape of the dataspace.
  74. mshape (read-only) => The shape of the selection region.
  75. Not guaranteed to fit within "shape", although
  76. the total number of points is less than
  77. product(shape).
  78. nselect (read-only) => Number of selected points. Always equal to
  79. product(mshape).
  80. broadcast(target_shape) => Return an iterable which yields dataspaces
  81. for read, based on target_shape.
  82. The base class represents "unshaped" selections (1-D).
  83. """
  84. def __init__(self, shape, spaceid=None):
  85. """ Create a selection. Shape may be None if spaceid is given. """
  86. if spaceid is not None:
  87. self._id = spaceid
  88. self._shape = spaceid.shape
  89. else:
  90. shape = tuple(shape)
  91. self._shape = shape
  92. self._id = h5s.create_simple(shape, (h5s.UNLIMITED,)*len(shape))
  93. self._id.select_all()
  94. @property
  95. def id(self):
  96. """ SpaceID instance """
  97. return self._id
  98. @property
  99. def shape(self):
  100. """ Shape of whole dataspace """
  101. return self._shape
  102. @property
  103. def nselect(self):
  104. """ Number of elements currently selected """
  105. return self._id.get_select_npoints()
  106. @property
  107. def mshape(self):
  108. """ Shape of selection (always 1-D for this class) """
  109. return (self.nselect,)
  110. @property
  111. def array_shape(self):
  112. """Shape of array to read/write (always 1-D for this class)"""
  113. return self.mshape
  114. # expand_shape and broadcast only really make sense for SimpleSelection
  115. def expand_shape(self, source_shape):
  116. if product(source_shape) != self.nselect:
  117. raise TypeError("Broadcasting is not supported for point-wise selections")
  118. return source_shape
  119. def broadcast(self, source_shape):
  120. """ Get an iterable for broadcasting """
  121. if product(source_shape) != self.nselect:
  122. raise TypeError("Broadcasting is not supported for point-wise selections")
  123. yield self._id
  124. def __getitem__(self, args):
  125. raise NotImplementedError("This class does not support indexing")
  126. class PointSelection(Selection):
  127. """
  128. Represents a point-wise selection. You can supply sequences of
  129. points to the three methods append(), prepend() and set(), or
  130. instantiate it with a single boolean array using from_mask().
  131. """
  132. def __init__(self, shape, spaceid=None, points=None):
  133. super().__init__(shape, spaceid)
  134. if points is not None:
  135. self._perform_selection(points, h5s.SELECT_SET)
  136. def _perform_selection(self, points, op):
  137. """ Internal method which actually performs the selection """
  138. points = np.asarray(points, order='C', dtype='u8')
  139. if len(points.shape) == 1:
  140. points.shape = (1,points.shape[0])
  141. if self._id.get_select_type() != h5s.SEL_POINTS:
  142. op = h5s.SELECT_SET
  143. if len(points) == 0:
  144. self._id.select_none()
  145. else:
  146. self._id.select_elements(points, op)
  147. @classmethod
  148. def from_mask(cls, mask, spaceid=None):
  149. """Create a point-wise selection from a NumPy boolean array """
  150. if not (isinstance(mask, np.ndarray) and mask.dtype.kind == 'b'):
  151. raise TypeError("PointSelection.from_mask only works with bool arrays")
  152. points = np.transpose(mask.nonzero())
  153. return cls(mask.shape, spaceid, points=points)
  154. def append(self, points):
  155. """ Add the sequence of points to the end of the current selection """
  156. self._perform_selection(points, h5s.SELECT_APPEND)
  157. def prepend(self, points):
  158. """ Add the sequence of points to the beginning of the current selection """
  159. self._perform_selection(points, h5s.SELECT_PREPEND)
  160. def set(self, points):
  161. """ Replace the current selection with the given sequence of points"""
  162. self._perform_selection(points, h5s.SELECT_SET)
  163. class SimpleSelection(Selection):
  164. """ A single "rectangular" (regular) selection composed of only slices
  165. and integer arguments. Can participate in broadcasting.
  166. """
  167. @property
  168. def mshape(self):
  169. """ Shape of current selection """
  170. return self._sel[1]
  171. @property
  172. def array_shape(self):
  173. scalar = self._sel[3]
  174. return tuple(x for x, s in zip(self.mshape, scalar) if not s)
  175. def __init__(self, shape, spaceid=None, hyperslab=None):
  176. super().__init__(shape, spaceid)
  177. if hyperslab is not None:
  178. self._sel = hyperslab
  179. else:
  180. # No hyperslab specified - select all
  181. rank = len(self.shape)
  182. self._sel = ((0,)*rank, self.shape, (1,)*rank, (False,)*rank)
  183. def expand_shape(self, source_shape):
  184. """Match the dimensions of an array to be broadcast to the selection
  185. The returned shape describes an array of the same size as the input
  186. shape, but its dimensions
  187. E.g. with a dataset shape (10, 5, 4, 2), writing like this::
  188. ds[..., 0] = np.ones((5, 4))
  189. The source shape (5, 4) will expand to (1, 5, 4, 1).
  190. Then the broadcast method below repeats that chunk 10
  191. times to write to an effective shape of (10, 5, 4, 1).
  192. """
  193. start, count, step, scalar = self._sel
  194. rank = len(count)
  195. remaining_src_dims = list(source_shape)
  196. eshape = []
  197. for idx in range(1, rank + 1):
  198. if len(remaining_src_dims) == 0 or scalar[-idx]: # Skip scalar axes
  199. eshape.append(1)
  200. else:
  201. t = remaining_src_dims.pop()
  202. if t == 1 or count[-idx] == t:
  203. eshape.append(t)
  204. else:
  205. raise TypeError("Can't broadcast %s -> %s" % (source_shape, self.array_shape)) # array shape
  206. if any([n > 1 for n in remaining_src_dims]):
  207. # All dimensions from target_shape should either have been popped
  208. # to match the selection shape, or be 1.
  209. raise TypeError("Can't broadcast %s -> %s" % (source_shape, self.array_shape)) # array shape
  210. # We have built eshape backwards, so now reverse it
  211. return tuple(eshape[::-1])
  212. def broadcast(self, source_shape):
  213. """ Return an iterator over target dataspaces for broadcasting.
  214. Follows the standard NumPy broadcasting rules against the current
  215. selection shape (self.mshape).
  216. """
  217. if self.shape == ():
  218. if product(source_shape) != 1:
  219. raise TypeError("Can't broadcast %s to scalar" % source_shape)
  220. self._id.select_all()
  221. yield self._id
  222. return
  223. start, count, step, scalar = self._sel
  224. rank = len(count)
  225. tshape = self.expand_shape(source_shape)
  226. chunks = tuple(x//y for x, y in zip(count, tshape))
  227. nchunks = product(chunks)
  228. if nchunks == 1:
  229. yield self._id
  230. else:
  231. sid = self._id.copy()
  232. sid.select_hyperslab((0,)*rank, tshape, step)
  233. for idx in range(nchunks):
  234. offset = tuple(x*y*z + s for x, y, z, s in zip(np.unravel_index(idx, chunks), tshape, step, start))
  235. sid.offset_simple(offset)
  236. yield sid
  237. class FancySelection(Selection):
  238. """
  239. Implements advanced NumPy-style selection operations in addition to
  240. the standard slice-and-int behavior.
  241. Indexing arguments may be ints, slices, lists of indices, or
  242. per-axis (1D) boolean arrays.
  243. Broadcasting is not supported for these selections.
  244. """
  245. @property
  246. def mshape(self):
  247. return self._mshape
  248. @property
  249. def array_shape(self):
  250. return self._array_shape
  251. def __init__(self, shape, spaceid=None, mshape=None, array_shape=None):
  252. super().__init__(shape, spaceid)
  253. if mshape is None:
  254. mshape = self.shape
  255. if array_shape is None:
  256. array_shape = mshape
  257. self._mshape = mshape
  258. self._array_shape = array_shape
  259. def expand_shape(self, source_shape):
  260. if not source_shape == self.array_shape:
  261. raise TypeError("Broadcasting is not supported for complex selections")
  262. return source_shape
  263. def broadcast(self, source_shape):
  264. if not source_shape == self.array_shape:
  265. raise TypeError("Broadcasting is not supported for complex selections")
  266. yield self._id
  267. def guess_shape(sid):
  268. """ Given a dataspace, try to deduce the shape of the selection.
  269. Returns one of:
  270. * A tuple with the selection shape, same length as the dataspace
  271. * A 1D selection shape for point-based and multiple-hyperslab selections
  272. * None, for unselected scalars and for NULL dataspaces
  273. """
  274. sel_class = sid.get_simple_extent_type() # Dataspace class
  275. sel_type = sid.get_select_type() # Flavor of selection in use
  276. if sel_class == h5s.NULL:
  277. # NULL dataspaces don't support selections
  278. return None
  279. elif sel_class == h5s.SCALAR:
  280. # NumPy has no way of expressing empty 0-rank selections, so we use None
  281. if sel_type == h5s.SEL_NONE: return None
  282. if sel_type == h5s.SEL_ALL: return tuple()
  283. elif sel_class != h5s.SIMPLE:
  284. raise TypeError("Unrecognized dataspace class %s" % sel_class)
  285. # We have a "simple" (rank >= 1) dataspace
  286. N = sid.get_select_npoints()
  287. rank = len(sid.shape)
  288. if sel_type == h5s.SEL_NONE:
  289. return (0,)*rank
  290. elif sel_type == h5s.SEL_ALL:
  291. return sid.shape
  292. elif sel_type == h5s.SEL_POINTS:
  293. # Like NumPy, point-based selections yield 1D arrays regardless of
  294. # the dataspace rank
  295. return (N,)
  296. elif sel_type != h5s.SEL_HYPERSLABS:
  297. raise TypeError("Unrecognized selection method %s" % sel_type)
  298. # We have a hyperslab-based selection
  299. if N == 0:
  300. return (0,)*rank
  301. bottomcorner, topcorner = (np.array(x) for x in sid.get_select_bounds())
  302. # Shape of full selection box
  303. boxshape = topcorner - bottomcorner + np.ones((rank,))
  304. def get_n_axis(sid, axis):
  305. """ Determine the number of elements selected along a particular axis.
  306. To do this, we "mask off" the axis by making a hyperslab selection
  307. which leaves only the first point along the axis. For a 2D dataset
  308. with selection box shape (X, Y), for axis 1, this would leave a
  309. selection of shape (X, 1). We count the number of points N_leftover
  310. remaining in the selection and compute the axis selection length by
  311. N_axis = N/N_leftover.
  312. """
  313. if(boxshape[axis]) == 1:
  314. return 1
  315. start = bottomcorner.copy()
  316. start[axis] += 1
  317. count = boxshape.copy()
  318. count[axis] -= 1
  319. # Throw away all points along this axis
  320. masked_sid = sid.copy()
  321. masked_sid.select_hyperslab(tuple(start), tuple(count), op=h5s.SELECT_NOTB)
  322. N_leftover = masked_sid.get_select_npoints()
  323. return N//N_leftover
  324. shape = tuple(get_n_axis(sid, x) for x in range(rank))
  325. if np.product(shape) != N:
  326. # This means multiple hyperslab selections are in effect,
  327. # so we fall back to a 1D shape
  328. return (N,)
  329. return shape