def test_spatiotemporal_cubes():
    def check_patch_coverage(files):
        rng = numpy.random.RandomState(1)
        inputs = [(name, array.shape) for name, array in six.iteritems(files)]
        shape = (5, 7, 7)
        for fname, index in spatiotemporal_cubes(inputs, shape, 50000, rng):
            cube = files[fname][index]
            if len(files[fname].shape) == 3:
                assert cube.shape == shape
            else:
                assert cube.shape[:3] == shape[:3]
            cube[...] = True
        for fname, array in six.iteritems(files):
            assert array.all()

    files = OrderedDict(
        file1=numpy.zeros((10, 30, 21), dtype=bool),
        file2=numpy.zeros((15, 25, 28), dtype=bool),
        file3=numpy.zeros((7, 18, 22), dtype=bool),
    )
    check_patch_coverage(files)

    # Check that stuff still works with an extra color channel dimension.
    files = OrderedDict(
        file1=numpy.zeros((10, 30, 21, 3), dtype=bool),
        file2=numpy.zeros((15, 25, 28, 3), dtype=bool),
        file3=numpy.zeros((7, 18, 22, 3), dtype=bool),
    )
    check_patch_coverage(files)
dataset/ETHPy150Open lisa-lab/pylearn2/pylearn2/utils/tests/test_video.py/test_spatiotemporal_cubes/original