You could subclass np.s_
to support "vectorized" slice notation:
import numpy as np
import types
class Block_Slice(np.lib.index_tricks.IndexExpression):
@staticmethod
def expand_slice(slice_, result_type=types.GeneratorType, force_type=False):
bc = np.broadcast(slice_.start, slice_.stop, slice_.step)
if bc.size == 1 and not force_type:
return slice_
if result_type in ('generator', types.GeneratorType):
return (slice(*idx) for idx in bc)
result_type = tuple if result_type=='tuple' else result_type
return result_type(slice(*idx) for idx in bc)
def __getitem__(self, item):
if isinstance(item, slice):
item = self.expand_slice(item, tuple)
elif isinstance(item, tuple):
item = tuple(j for i in item for j in
(self.expand_slice(i, force_type=True)
if isinstance(i, slice) else (i,)))
return super().__getitem__(item)
bs_ = Block_Slice(maketuple=False)
a = sum(np.ogrid[:2000:1000, :300:100, :40:10, :6])
left = [1, 0, 0]
right = [2, 2, 4]
a[bs_[left:right]]
# array([[[[1000, 1001, 1002, 1003, 1004, 1005],
# [1010, 1011, 1012, 1013, 1014, 1015],
# [1020, 1021, 1022, 1023, 1024, 1025],
# [1030, 1031, 1032, 1033, 1034, 1035]],
# [[1100, 1101, 1102, 1103, 1104, 1105],
# [1110, 1111, 1112, 1113, 1114, 1115],
# [1120, 1121, 1122, 1123, 1124, 1125],
# [1130, 1131, 1132, 1133, 1134, 1135]]]])
a[bs_[..., range(3):5]]
# array([[[[ 12, 13, 14],
# [ 22, 23, 24],
# [ 32, 33, 34]],
# [[ 112, 113, 114],
# [ 122, 123, 124],
# [ 132, 133, 134]],
# [[ 212, 213, 214],
# [ 222, 223, 224],
# [ 232, 233, 234]]],
# [[[1012, 1013, 1014],
# [1022, 1023, 1024],
# [1032, 1033, 1034]],
# [[1112, 1113, 1114],
# [1122, 1123, 1124],
# [1132, 1133, 1134]],
# [[1212, 1213, 1214],
# [1222, 1223, 1224],
# [1232, 1233, 1234]]]])
a[bs_[:right, [4, 2, 1]]]
# array([[[[ 4, 2, 1],
# [ 14, 12, 11],
# [ 24, 22, 21],
# [ 34, 32, 31]],
# [[ 104, 102, 101],
# [ 114, 112, 111],
# [ 124, 122, 121],
# [ 134, 132, 131]]],
# [[[1004, 1002, 1001],
# [1014, 1012, 1011],
# [1024, 1022, 1021],
# [1034, 1032, 1031]],
# [[1104, 1102, 1101],
# [1114, 1112, 1111],
# [1124, 1122, 1121],
# [1134, 1132, 1131]]]])