def test_get_bricks_from_slice_1d(self):
        total_domain = (15, )
        brick_extents = (((0, 4), ), ((5, 9), ), ((10, 14), ))
        rtree_extents = ((0, 0, 4, 0), (5, 0, 9, 0), (10, 0, 14, 0))

        brick_0 = (0, ((0, 4), ))
        brick_1 = (1, ((5, 9), ))
        brick_2 = (2, ((10, 14), ))

        from coverage_model.test.bricking_assessment_utility import BrickingAssessor
        rtree = RTreeProxy()
        for x in BrickingAssessor.rtree_populator(rtree_extents,
                                                  brick_extents):
            rtree.insert(*x)

        # Try a variety of slices
        self._get_bricks_assert(slice(None), rtree, total_domain, 3,
                                [brick_0, brick_1, brick_2])

        self._get_bricks_assert(slice(None, None, 3), rtree, total_domain, 3,
                                [brick_0, brick_1, brick_2])

        self._get_bricks_assert(slice(0, 3), rtree, total_domain, 1, [brick_0])

        self._get_bricks_assert(slice(5, 9), rtree, total_domain, 1, [brick_1])

        self._get_bricks_assert(slice(6, None), rtree, total_domain, 2,
                                [brick_1, brick_2])

        self._get_bricks_assert(
            slice(None, None, 10), rtree, total_domain, 3,
            [brick_0, brick_1, brick_2
             ])  # three bricks, tho the middle one isn't needed

        self._get_bricks_assert(([1, 3], ), rtree, total_domain, 1, [brick_0])

        self._get_bricks_assert(([2, 4, 7], ), rtree, total_domain, 2,
                                [brick_0, brick_1])

        self._get_bricks_assert(
            ([3, 12], ), rtree, total_domain, 3,
            [brick_0, brick_1, brick_2
             ])  # three bricks, tho the middle one isn't needed

        self._get_bricks_assert(1, rtree, total_domain, 1, [brick_0])

        self._get_bricks_assert(6, rtree, total_domain, 1, [brick_1])

        self._get_bricks_assert(13, rtree, total_domain, 1, [brick_2])
    def test_get_bricks_from_slice_3d(self):
        total_domain = (10, 15, 5)
        brick_extents = (((0, 4), (0, 4), (0, 4)), ((0, 4), (5, 9), (0, 4)),
                         ((0, 4), (10, 14), (0, 4)), ((5, 9), (0, 4), (0, 4)),
                         ((5, 9), (5, 9), (0, 4)), ((5, 9), (10, 14), (0, 4)))
        rtree_extents = ((0, 0, 0, 4, 4, 4), (0, 5, 0, 4, 9, 4),
                         (0, 10, 0, 4, 14, 4), (5, 0, 0, 9, 4, 4),
                         (5, 5, 0, 9, 9, 4), (5, 10, 0, 9, 14, 4))

        brick_0 = (0, ((0, 4), (0, 4), (0, 4)))
        brick_1 = (1, ((0, 4), (5, 9), (0, 4)))
        brick_2 = (2, ((0, 4), (10, 14), (0, 4)))
        brick_3 = (3, ((5, 9), (0, 4), (0, 4)))
        brick_4 = (4, ((5, 9), (5, 9), (0, 4)))
        brick_5 = (5, ((5, 9), (10, 14), (0, 4)))

        from coverage_model.test.bricking_assessment_utility import BrickingAssessor
        rtree = RTreeProxy()
        for x in BrickingAssessor.rtree_populator(rtree_extents,
                                                  brick_extents):
            rtree.insert(*x)

        # Get all bricks
        self._get_bricks_assert(
            (slice(None), ) * 3, rtree, total_domain, 6,
            [brick_0, brick_1, brick_2, brick_3, brick_4, brick_5])

        self._get_bricks_assert((0, 0, 0), rtree, total_domain, 1, [brick_0])

        self._get_bricks_assert((8, 5, 2), rtree, total_domain, 1, [brick_4])

        self._get_bricks_assert((4, 12, 1), rtree, total_domain, 1, [brick_2])

        self._get_bricks_assert((9, 13, [0, 2]), rtree, total_domain, 1,
                                [brick_5])

        self._get_bricks_assert((8, [3, 5, 12], 0), rtree, total_domain, 3,
                                [brick_3, brick_4, brick_5])

        self._get_bricks_assert(([5, 9], 10, 0), rtree, total_domain, 1,
                                [brick_5])

        self._get_bricks_assert(([5, 9], [4, 12, 13], 0), rtree, total_domain,
                                3, [brick_3, brick_4, brick_5])

        self._get_bricks_assert(([2, 4], [2, 11], [1, 3, 4]), rtree,
                                total_domain, 3, [brick_0, brick_1, brick_2])

        self._get_bricks_assert(([2, 3, 9], 12, [1, 3, 4]), rtree,
                                total_domain, 2, [brick_2, brick_5])

        self._get_bricks_assert((slice(None), 12, [1, 3, 4]), rtree,
                                total_domain, 2, [brick_2, brick_5])

        self._get_bricks_assert((slice(1, 7), 3, [1, 3, 4]), rtree,
                                total_domain, 2, [brick_0, brick_3])

        self._get_bricks_assert((slice(3, 4), 7, [1, 3, 4]), rtree,
                                total_domain, 1, [brick_1])

        self._get_bricks_assert(
            (slice(2, 8, 7), [1, 6, 12], 4), rtree, total_domain, 6,
            [brick_0, brick_1, brick_2, brick_3, brick_4, brick_5])

        self._get_bricks_assert((slice(2, 4, 7), slice(None), 2), rtree,
                                total_domain, 3, [brick_0, brick_1, brick_2])

        self._get_bricks_assert(
            (slice(None, 4), slice(9, None, 2), slice(None)), rtree,
            total_domain, 2, [brick_1, brick_2])

        self._get_bricks_assert(
            (slice(None, 6, 4), slice(12, None, 2), slice(3, None)), rtree,
            total_domain, 2, [brick_2, brick_5])

        self._get_bricks_assert(
            (slice(None, 8), slice(6, 13, 4), slice(None, None, 3)), rtree,
            total_domain, 4, [brick_1, brick_2, brick_4, brick_5])
    def test_get_bricks_from_slice_2d(self):
        total_domain = (15, 10)
        brick_extents = (((0, 4), (0, 4)), ((0, 4), (5, 9)), ((5, 9), (0, 4)),
                         ((5, 9), (5, 9)), ((10, 14), (0, 4)), ((10, 14), (5,
                                                                           9)))
        rtree_extents = ((0, 0, 4, 4), (0, 5, 4, 9), (5, 0, 9, 4),
                         (5, 5, 9, 9), (10, 0, 14, 4), (10, 5, 14, 9))

        brick_0 = (0, ((0, 4), (0, 4)))
        brick_1 = (1, ((0, 4), (5, 9)))
        brick_2 = (2, ((5, 9), (0, 4)))
        brick_3 = (3, ((5, 9), (5, 9)))
        brick_4 = (4, ((10, 14), (0, 4)))
        brick_5 = (5, ((10, 14), (5, 9)))

        from coverage_model.test.bricking_assessment_utility import BrickingAssessor
        rtree = RTreeProxy()
        for x in BrickingAssessor.rtree_populator(rtree_extents,
                                                  brick_extents):
            rtree.insert(*x)

        # Get all bricks
        self._get_bricks_assert(
            (slice(None), ) * 2, rtree, total_domain, 6,
            [brick_0, brick_1, brick_2, brick_3, brick_4, brick_5])

        self._get_bricks_assert(
            (slice(None), slice(None, 8)), rtree, total_domain, 6,
            [brick_0, brick_1, brick_2, brick_3, brick_4, brick_5])

        self._get_bricks_assert((slice(None), slice(None, 4)), rtree,
                                total_domain, 3, [brick_0, brick_2, brick_4])

        self._get_bricks_assert((slice(7, 12), slice(5, 8)), rtree,
                                total_domain, 2, [brick_3, brick_5])

        self._get_bricks_assert(
            (slice(2, 14, 3), slice(2, 7)), rtree, total_domain, 6,
            [brick_0, brick_1, brick_2, brick_3, brick_4, brick_5])

        self._get_bricks_assert(
            (slice(2, 14, 10), slice(2, 7)), rtree, total_domain, 6,
            [brick_0, brick_1, brick_2, brick_3, brick_4, brick_5])

        self._get_bricks_assert((0, slice(2, 8, 3)), rtree, total_domain, 2,
                                [brick_0, brick_1])

        self._get_bricks_assert((6, slice(2, 7)), rtree, total_domain, 2,
                                [brick_2, brick_3])

        self._get_bricks_assert((slice(None, 12), 7), rtree, total_domain, 3,
                                [brick_1, brick_3, brick_5])

        self._get_bricks_assert((12, slice(2, None, 4)), rtree, total_domain,
                                2, [brick_4, brick_5])

        self._get_bricks_assert(([1, 2], 9), rtree, total_domain, 1, [brick_1])

        self._get_bricks_assert(([0, 14], 3), rtree, total_domain, 3,
                                [brick_0, brick_2, brick_4])

        self._get_bricks_assert((3, [1, 8]), rtree, total_domain, 2,
                                [brick_0, brick_1])

        self._get_bricks_assert(([2, 5], [1, 8]), rtree, total_domain, 4,
                                [brick_0, brick_1, brick_2, brick_3])

        self._get_bricks_assert(([6, 9], [1, 8]), rtree, total_domain, 2,
                                [brick_2, brick_3])

        self._get_bricks_assert(([2, 8, 13], [7, 8]), rtree, total_domain, 3,
                                [brick_1, brick_3, brick_5])
class BrickingAssessor(object):
    def __init__(self,
                 total_domain=(10, 10),
                 brick_size=5,
                 use_hdf=False,
                 root_dir='test_data/multi_dim_trials',
                 guid=None,
                 dtype='int16'):
        self.total_domain = total_domain
        self.brick_sizes = tuple(brick_size for x in total_domain)
        self.use_hdf = use_hdf
        self.dtype = np.dtype(dtype).name
        if self.use_hdf:
            self.guid = guid or create_guid()
            name = '%s_%s' % (self.guid, self.dtype)
            self.root_dir = root_dir
            if not os.path.exists(self.root_dir):
                os.makedirs(self.root_dir)

            if os.path.exists(os.path.join(self.root_dir, name)):
                shutil.rmtree(os.path.join(self.root_dir, name))

            self.master_manager = MasterManager(
                self.root_dir, name, name='md_test_{0}'.format(name))

            self.master_manager.flush()

            pc = ParameterContext('test_param',
                                  param_type=QuantityType(self.dtype),
                                  fill_value=-1)
            self.param_manager = ParameterManager(
                os.path.join(self.root_dir, name, pc.name), pc.name)
            self.param_manager.parameter_context = pc
            self.master_manager.create_group(pc.name)

            self.param_manager.flush()

        self.bricks = {}

        self.brick_origins = bricking_utils.calc_brick_origins(
            self.total_domain, self.brick_sizes)
        self.brick_extents, self.rtree_extents = bricking_utils.calc_brick_and_rtree_extents(
            self.brick_origins, self.brick_sizes)
        self.build_bricks()

        self.rtree = RTreeProxy()
        for x in BrickingAssessor.rtree_populator(self.rtree_extents,
                                                  self.brick_extents):
            self.rtree.insert(*x)

    @classmethod
    def rtree_populator(cls, rtree_extents, brick_extents):
        for i, e in enumerate(rtree_extents):
            yield i, e, brick_extents[i]

    def _get_numpy_array(self, shape):
        if not isinstance(shape, tuple):
            shape = tuple(shape)

        return np.arange(utils.prod(shape), dtype=self.dtype).reshape(shape)

    def build_bricks(self):
        for x in xrange(len(self.brick_origins)):
            if not self.use_hdf:
                self.bricks[x] = np.empty(self.brick_sizes, dtype=self.dtype)
                self.bricks[x].fill(-1)
            else:
                id = str(x)
                fn = '{0}.hdf5'.format(id)
                pth = os.path.join(self.param_manager.root_dir, fn)
                relpth = os.path.join(
                    self.param_manager.root_dir.replace(
                        self.master_manager.root_dir, '.'), fn)
                lnpth = '/{0}/{1}'.format(self.param_manager.parameter_name,
                                          id)

                self.master_manager.add_external_link(lnpth, relpth, id)
                self.bricks[x] = pth

    def reset_bricks(self):
        for i, arr in enumerate(self.bricks.itervalues()):
            if not self.use_hdf:
                arr.fill(-1)
            else:
                with h5py.File(arr) as f:
                    ds = f.require_dataset(str(i),
                                           shape=self.brick_sizes,
                                           dtype=self.dtype,
                                           chunks=None,
                                           fillvalue=-1)
                    ds[:] = -1

    def put_values_to_bricks(self, slice_, values):
        slice_ = utils.fix_slice(slice_, self.total_domain)
        bricks = bricking_utils.get_bricks_from_slice(
            slice_, self.rtree, self.total_domain
        )  # this is a list of tuples [(b_id, (bounds...),), ...]

        values = np.asanyarray(values)
        v_shp = values.shape
        log.debug('value_shape: %s', v_shp)
        s_shp = utils.slice_shape(slice_, self.total_domain)
        log.debug('slice_shape: %s', s_shp)
        is_broadcast = False
        if v_shp == ():
            log.debug('Broadcast!!')
            is_broadcast = True
            value_slice = ()
        elif v_shp != s_shp:
            if v_shp == tuple([
                    i for i in s_shp if i != 1
            ]):  # Missing dimensions are singleton, just reshape to fit
                values = values.reshape(s_shp)
                v_shp = values.shape
            else:
                raise IndexError(
                    'Shape of \'value\' is not compatible with \'slice_\': slice_ shp == {0}\tvalue shp == {1}'
                    .format(s_shp, v_shp))
        else:
            value_slice = None

        log.debug('value_shape: %s', v_shp)

        for b in bricks:
            # b is (brick_id, (brick_bounds per dim...),)
            bid, bbnds = b
            log.debug('Determining slice for brick: %s', b)
            bexts = tuple([x + 1 for x in zip(*bbnds)[1]
                           ])  # Shift from index to size
            log.debug('bid=%s, bbnds=%s, bexts=%s', bid, bbnds, bexts)

            brick_slice, brick_mm = bricking_utils.get_brick_slice_nd(
                slice_, bbnds)

            if None in brick_slice:  # Brick does not contain any of the requested indices
                log.debug(
                    'Brick does not contain any of the requested indices: Move to next brick'
                )
                continue

            try:
                brick_slice = utils.fix_slice(brick_slice, bexts)
            except IndexError:
                log.debug('Malformed brick_slice: move to next brick')
                continue

            if not is_broadcast:
                value_slice = bricking_utils.get_value_slice_nd(
                    slice_, v_shp, bbnds, brick_slice, brick_mm)

                try:
                    value_slice = utils.fix_slice(value_slice, v_shp)
                except IndexError:
                    log.debug('Malformed value_slice: move to next brick')
                    continue

            log.debug(
                '\nbrick %s:\n\tbrick_slice %s=%s\n\tmin/max=%s\n\tvalue_slice %s=%s',
                b, utils.slice_shape(brick_slice,
                                     bexts), brick_slice, brick_mm,
                utils.slice_shape(value_slice, v_shp), value_slice)
            v = values[value_slice]
            log.debug('\nvalues %s=\n%s', v.shape, v)
            if not self.use_hdf:
                self.bricks[bid][brick_slice] = v
            else:
                fi = self.bricks[bid]
                with h5py.File(fi) as f:
                    ds = f.require_dataset(str(bid),
                                           shape=self.brick_sizes,
                                           dtype=self.dtype,
                                           chunks=None,
                                           fillvalue=-1)
                    ds[brick_slice] = v

    def get_values_from_bricks(self, slice_):
        slice_ = utils.fix_slice(slice_, self.total_domain)
        bricks = bricking_utils.get_bricks_from_slice(
            slice_, self.rtree, self.total_domain
        )  # this is a list of tuples [(b_id, (bounds...),), ...]

        ret_shp = utils.slice_shape(slice_, self.total_domain)
        ret_arr = np.empty(ret_shp, dtype=self.dtype)

        for b in bricks:
            bid, bbnds = b
            brick_slice, brick_mm = bricking_utils.get_brick_slice_nd(
                slice_, bbnds)

            if None in brick_slice:
                continue

            ret_slice = bricking_utils.get_value_slice_nd(
                slice_, ret_shp, bbnds, brick_slice, brick_mm)

            if not self.use_hdf:
                ret_vals = self.bricks[bid][brick_slice]
            else:
                fi = self.bricks[bid]
                with h5py.File(fi) as f:
                    ds = f.require_dataset(str(bid),
                                           shape=self.brick_sizes,
                                           dtype=self.dtype,
                                           chunks=None,
                                           fillvalue=-1)
                    ret_vals = ds[brick_slice]

            ret_arr[ret_slice] = ret_vals

        ret_arr = ret_arr.squeeze()

        if ret_arr.size == 1:
            if ret_arr.ndim == 0:
                ret_arr = ret_arr[()]
            else:
                ret_arr = ret_arr[0]

        return ret_arr
class BrickingAssessor(object):
    def __init__(self, total_domain=(10, 10), brick_size=5, use_hdf=False, root_dir='test_data/multi_dim_trials',
                 guid=None, dtype='int16'):
        self.total_domain = total_domain
        self.brick_sizes = tuple(brick_size for x in total_domain)
        self.use_hdf = use_hdf
        self.dtype = np.dtype(dtype).name
        if self.use_hdf:
            self.guid = guid or create_guid()
            name = '%s_%s' % (self.guid, self.dtype)
            self.root_dir = root_dir
            if not os.path.exists(self.root_dir):
                os.makedirs(self.root_dir)

            if os.path.exists(os.path.join(self.root_dir, name)):
                shutil.rmtree(os.path.join(self.root_dir, name))

            self.master_manager = MasterManager(self.root_dir, name, name='md_test_{0}'.format(name))

            self.master_manager.flush()

            pc = ParameterContext('test_param', param_type=QuantityType(self.dtype), fill_value=-1)
            self.param_manager = ParameterManager(os.path.join(self.root_dir, name, pc.name), pc.name)
            self.param_manager.parameter_context = pc
            self.master_manager.create_group(pc.name)

            self.param_manager.flush()

        self.bricks = {}

        self.brick_origins = bricking_utils.calc_brick_origins(self.total_domain, self.brick_sizes)
        self.brick_extents, self.rtree_extents = bricking_utils.calc_brick_and_rtree_extents(self.brick_origins,
                                                                                             self.brick_sizes)
        self.build_bricks()

        self.rtree = RTreeProxy()
        for x in BrickingAssessor.rtree_populator(self.rtree_extents, self.brick_extents):
            self.rtree.insert(*x)

    @classmethod
    def rtree_populator(cls, rtree_extents, brick_extents):
        for i, e in enumerate(rtree_extents):
            yield i, e, brick_extents[i]

    def _get_numpy_array(self, shape):
        if not isinstance(shape, tuple):
            shape = tuple(shape)

        return np.arange(utils.prod(shape), dtype=self.dtype).reshape(shape)

    def build_bricks(self):
        for x in xrange(len(self.brick_origins)):
            if not self.use_hdf:
                self.bricks[x] = np.empty(self.brick_sizes, dtype=self.dtype)
                self.bricks[x].fill(-1)
            else:
                id = str(x)
                fn = '{0}.hdf5'.format(id)
                pth = os.path.join(self.param_manager.root_dir, fn)
                relpth = os.path.join(self.param_manager.root_dir.replace(self.master_manager.root_dir, '.'), fn)
                lnpth = '/{0}/{1}'.format(self.param_manager.parameter_name, id)

                self.master_manager.add_external_link(lnpth, relpth, id)
                self.bricks[x] = pth

    def reset_bricks(self):
        for i, arr in enumerate(self.bricks.itervalues()):
            if not self.use_hdf:
                arr.fill(-1)
            else:
                with HDFLockingFile(arr, mode='a') as f:
                    ds = f.require_dataset(str(i), shape=self.brick_sizes, dtype=self.dtype, chunks=None, fillvalue=-1)
                    ds[:] = -1

    def put_values_to_bricks(self, slice_, values):
        slice_ = utils.fix_slice(slice_, self.total_domain)
        bricks = bricking_utils.get_bricks_from_slice(slice_, self.rtree,
                                                      self.total_domain) # this is a list of tuples [(b_id, (bounds...),), ...]

        values = np.asanyarray(values)
        v_shp = values.shape
        log.debug('value_shape: %s', v_shp)
        s_shp = utils.slice_shape(slice_, self.total_domain)
        log.debug('slice_shape: %s', s_shp)
        is_broadcast = False
        if v_shp == ():
            log.debug('Broadcast!!')
            is_broadcast = True
            value_slice = ()
        elif v_shp != s_shp:
            if v_shp == tuple([i for i in s_shp if i != 1]): # Missing dimensions are singleton, just reshape to fit
                values = values.reshape(s_shp)
                v_shp = values.shape
            else:
                raise IndexError(
                    'Shape of \'value\' is not compatible with \'slice_\': slice_ shp == {0}\tvalue shp == {1}'.format(
                        s_shp, v_shp))
        else:
            value_slice = None

        log.debug('value_shape: %s', v_shp)

        for b in bricks:
            # b is (brick_id, (brick_bounds per dim...),)
            bid, bbnds = b
            log.debug('Determining slice for brick: %s', b)
            bexts = tuple([x + 1 for x in zip(*bbnds)[1]]) # Shift from index to size
            log.debug('bid=%s, bbnds=%s, bexts=%s', bid, bbnds, bexts)

            brick_slice, brick_mm = bricking_utils.get_brick_slice_nd(slice_, bbnds)

            if None in brick_slice: # Brick does not contain any of the requested indices
                log.debug('Brick does not contain any of the requested indices: Move to next brick')
                continue

            try:
                brick_slice = utils.fix_slice(brick_slice, bexts)
            except IndexError:
                log.debug('Malformed brick_slice: move to next brick')
                continue

            if not is_broadcast:
                value_slice = bricking_utils.get_value_slice_nd(slice_, v_shp, bbnds, brick_slice, brick_mm)

                try:
                    value_slice = utils.fix_slice(value_slice, v_shp)
                except IndexError:
                    log.debug('Malformed value_slice: move to next brick')
                    continue

            log.debug('\nbrick %s:\n\tbrick_slice %s=%s\n\tmin/max=%s\n\tvalue_slice %s=%s', b,
                      utils.slice_shape(brick_slice, bexts), brick_slice, brick_mm,
                      utils.slice_shape(value_slice, v_shp), value_slice)
            v = values[value_slice]
            log.debug('\nvalues %s=\n%s', v.shape, v)
            if not self.use_hdf:
                self.bricks[bid][brick_slice] = v
            else:
                fi = self.bricks[bid]
                with HDFLockingFile(fi, 'a') as f:
                    ds = f.require_dataset(str(bid), shape=self.brick_sizes, dtype=self.dtype, chunks=None,
                                           fillvalue=-1)
                    ds[brick_slice] = v

    def get_values_from_bricks(self, slice_):
        slice_ = utils.fix_slice(slice_, self.total_domain)
        bricks = bricking_utils.get_bricks_from_slice(slice_, self.rtree,
                                                      self.total_domain) # this is a list of tuples [(b_id, (bounds...),), ...]

        ret_shp = utils.slice_shape(slice_, self.total_domain)
        ret_arr = np.empty(ret_shp, dtype=self.dtype)

        for b in bricks:
            bid, bbnds = b
            brick_slice, brick_mm = bricking_utils.get_brick_slice_nd(slice_, bbnds)

            if None in brick_slice:
                continue

            ret_slice = bricking_utils.get_value_slice_nd(slice_, ret_shp, bbnds, brick_slice, brick_mm)

            if not self.use_hdf:
                ret_vals = self.bricks[bid][brick_slice]
            else:
                fi = self.bricks[bid]
                with HDFLockingFile(fi) as f:
                    ds = f.require_dataset(str(bid), shape=self.brick_sizes, dtype=self.dtype, chunks=None,
                                           fillvalue=-1)
                    ret_vals = ds[brick_slice]

            ret_arr[ret_slice] = ret_vals

        ret_arr = ret_arr.squeeze()

        if ret_arr.size == 1:
            if ret_arr.ndim == 0:
                ret_arr = ret_arr[()]
            else:
                ret_arr = ret_arr[0]

        return ret_arr