22156500136

Committed 18 Feb 2026 08:31PM UTC coverage: 55.908%. Remained the same

Build # 22156500136

Build Type

push

github

Committed by

mhasself

Commit Message

feat: skip shape check for RangesMatrix math operations

Coverage Stats

1 of 11 new or added lines in 1 file covered. (9.09%)

1339 of 2395 relevant lines covered (55.91%)

0.56 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

75.0

/python/proj/ranges.py

import so3g
import numpy as np

"""Objects will self report as being of type "RangesInt32" rather than
Ranges.  But let's try to use so3g.proj.Ranges when testing types and
making new ones and stuff."""

Ranges = so3g.RangesInt32


class RangesMatrix():
    """This is a wrapper for multi-dimensional grids of Ranges objects.
    This can be used to store Ranges objects per-detector (the 2-d
    case) or per-thread and per-detector (a 3-d case, required if
    using OpenMP).  The right-most axis always corresponds to the
    sample axis carried by the Ranges objects.

    In addition to .shape and multi-dimensional slicing, it supports
    inversion (the ~ operator) and multiplication against other
    RangesMatrix or Ranges objects, with broadcasting rules similar to
    standard numpy array rules.

    """
    def __init__(self, items=[], child_shape=None, skip_shape_check=False):
        self.ranges = [x for x in items]
        if len(items):
            child_shape = items[0].shape
            if not skip_shape_check:
                assert all([(item.shape == child_shape) for item in items])
        elif child_shape is None:
            child_shape = ()
        self._child_shape = child_shape

    def __repr__(self):
        return 'RangesMatrix(' + ','.join(map(str, self.shape)) + ')'

    def __len__(self):
        return self.shape[0]

    def copy(self):
        return RangesMatrix([x.copy() for x in self.ranges],
                            child_shape=self.shape[1:])

    def zeros_like(self):
        return RangesMatrix([x.zeros_like() for x in self.ranges],
                            child_shape=self.shape[1:])
    
    def ones_like(self):
        return RangesMatrix([x.ones_like() for x in self.ranges],
                            child_shape=self.shape[1:])

    def buffer(self, buff):
        [x.buffer(buff) for x in self.ranges]
        ## just to make this work like Ranges.buffer()
        return self
    
    def buffered(self, buff):
        out = self.copy()
        [x.buffer(buff) for x in out.ranges]
        return out

    @property
    def shape(self):
        if len(self.ranges) == 0:
            return (0,) + self._child_shape
        return (len(self.ranges),) + self.ranges[0].shape

    def __getitem__(self, index):
        """RangesMatrix supports multi-dimensional indexing, slicing, and
        numpy-style advanced indexing (with some restrictions).

        The right-most axis of RangesMatrix has special restrictions,
        since it corresponds to Ranges objects: it can only accept
        slice-based indexing, not integer or advanced indexing.

        To guarantee that you get copies, rather than references to,
        the lowest level Ranges objects, make sure the index tuple
        includes a slice along the final (right-most) axis.  For
        example::

          # Starting point
          rm = RangesMatrix.zeros(10, 10000)

          # Get and modify an element in rm ...
          r = rm[0]
          r.add_interval(0, 10)  # This _does_ affect rm!

          # Get a copy of an element from rm.
          r = rm[0,:]
          r.add_interval(0, 10)  # This will not affect rm.

          # More generally, this is equivalent to rm1 = rm.copy():
          new_rm = rm[..., :]

        """
        if not isinstance(index, tuple):
            index = (index,)

        eidx = [i for i, a in enumerate(index) if a is Ellipsis]
        if len(eidx) == 1:
            # Fill in missing slices.
            eidx = eidx[0]
            n_free = len(self.shape) - sum([e is not None for e in index]) + 1
            index = index[:eidx] + tuple([slice(None)] * n_free) + index[eidx+1:]
        elif len(eidx) > 1:
            raise IndexError("An index can only have a single ellipsis ('...')")

        return _gibasic(self, index)

    def __add__(self, x):
        if isinstance(x, Ranges):
            return self.__class__([d + x for d in self.ranges], skip_shape_check=True)
        elif isinstance(x, RangesMatrix):
            # Check for shape compatibility.
            nd_a = len(self.shape)
            nd_b = len(x.shape)
            ndim = min(nd_a, nd_b)
            ok = [(a==1 or b==1 or a == b) for a,b in zip(self.shape[-ndim:], x.shape[-ndim:])]
            if not all(ok) or nd_a < nd_b:
                raise ValueError('Operands have incompatible shapes: %s %s' %
                                 self.shape, x.shape)
            if nd_a == nd_b:
                # Broadcast if either has shape 1...
                if x.shape[0] == 1:
                    return self.__class__([r + x[0] for r in self.ranges], skip_shape_check=True)
                elif self.shape[0] == 1:
                    return self.__class__([self.ranges[0] + _x for _x in x], skip_shape_check=True)
                elif self.shape[0] == x.shape[0]:
                    return self.__class__([r + d for r, d in zip(self.ranges, x)], skip_shape_check=True)
            return self.__class__([r + x for r in self.ranges], skip_shape_check=True)
        
    def __mul__(self, x):
        if isinstance(x, Ranges):
            return self.__class__([d * x for d in self.ranges], skip_shape_check=True)
        elif isinstance(x, RangesMatrix):
            # Check for shape compatibility.
            nd_a = len(self.shape)
            nd_b = len(x.shape)
            ndim = min(nd_a, nd_b)
            ok = [(a==1 or b==1 or a == b) for a,b in zip(self.shape[-ndim:], x.shape[-ndim:])]
            if not all(ok) or nd_a < nd_b:
                raise ValueError('Operands have incompatible shapes: %s %s' %
                                 self.shape, x.shape)
            if nd_a == nd_b:
                # Broadcast if either has shape 1...
                if x.shape[0] == 1:
                    return self.__class__([r * x[0] for r in self.ranges], skip_shape_check=True)
                elif self.shape[0] == 1:
                    return self.__class__([self.ranges[0] * _x for _x in x], skip_shape_check=True)
                elif self.shape[0] == x.shape[0]:
                    return self.__class__([r * d for r, d in zip(self.ranges, x)], skip_shape_check=True)
            return self.__class__([r * x for r in self.ranges], skip_shape_check=True)

    def __invert__(self):
        return self.__class__([~x for x in self.ranges], skip_shape_check=True)

    @staticmethod
    def concatenate(items, axis=0):
        """Static method to concatenate multiple RangesMatrix (or Ranges)
        objects along the specified axis.

        """
        # Check shape compatibility...
        s = list(items[0].shape)
        s[axis] = -1
        for item in items[1:]:
            s1 = list(item.shape)
            s1[axis] = -1
            if s != s1:
                raise ValueError('Contributed items must have same shape on non-cat axis.')
        def collect(items, join_depth):
            # Recurse through items in order to concatenate along axis join_depth.
            ranges = []
            if join_depth > 0:
                for co_items in zip(*items):
                    ranges.append(collect(co_items, join_depth - 1))
            else:
                if isinstance(items[0], RangesMatrix):
                    for item in items:
                        ranges.extend(item.ranges)
                else:
                    # The items are Ranges, then.
                    n, r = 0, []
                    for item in items:
                        r.extend(item.ranges() + n)
                        n += item.count
                    r = Ranges.from_array(
                        np.array(r, dtype='int32').reshape((-1, 2)), n)
                    return r
            return RangesMatrix(ranges, child_shape=items[0].shape[1:])
        return collect(items, axis)

    def close_gaps(self, gap_size=0):
        """Call close_gaps(gap_size) on all children.  Any ranges that abutt
        each other within the gap_size are merged into a single entry.
        Usually a gap_size of 0 is not possible, but if a caller is
        carefully using append_interval_no_check, then it can happen.

        """
        for r in self.ranges:
            r.close_gaps(gap_size)

    def get_stats(self):
        samples, intervals = [], []
        for r in self.ranges:
            if isinstance(r, Ranges):
                ra = r.ranges()
                samples.append(np.dot(ra, [-1, 1]).sum())
                intervals.append(ra.shape[0])
            else:
                sub = r.get_stats()
                samples.append(sum(sub['samples']))
                intervals.append(sum(sub['intervals']))
        return {
            'samples': samples,
            'intervals': intervals}

    @staticmethod
    def full(shape, fill_value):
        """Construct a RangesMatrix with the specified shape, initialized to
        fill_value.

        Args:
          shape (tuple of int): The shape.  Must have at least one
            element.  If there is only one element, a Ranges object is
            returned, not a RangesMatrix.
          fill_value (bool): True or False.  If False, the wrapped
            Ranges objects will have no intervals.  If True, the
            wrapped Ranges objects will all have a single interval
            spanning their entire domain.

        Returns:
          Ranges object (if shape has a single element) or a
          RangesMatrix object (len(shape) > 1).

        See also: zeros, ones.

        """
        assert fill_value in [True, False]
        if isinstance(shape, int):
            shape = (shape,)
        assert(len(shape) > 0)
        if len(shape) == 1:
            r = Ranges(shape[0])
            if fill_value:
                r = ~r
            return r
        return RangesMatrix([RangesMatrix.full(shape[1:], fill_value)
                             for i in range(shape[0])],
                            child_shape=shape[1:])

    @classmethod
    def zeros(cls, shape):
        """Equivalent to full(shape, False)."""
        return cls.full(shape, False)

    @classmethod
    def ones(cls, shape):
        """Equivalent to full(shape, True)."""
        return cls.full(shape, True)

    @classmethod
    def from_mask(cls, mask):
        """Create a RangesMatrix from a boolean mask.  The input mask can have
        any dimensionality greater than 0 but be aware that if ndim==1
        then a Ranges object, and not a RangesMatrix, is returned.

        Args:
          mask (ndarray): Must be boolean array with at least 1 dimension.

        Returns:
          RangesMatrix with the same shape as mask, with ranges
          corresponding to the intervals where mask was True.

        """
        assert(mask.ndim > 0)
        if mask.ndim == 1:
           return Ranges.from_mask(mask)
        if len(mask) == 0:
            return cls(child_shape=mask.shape[1:])
        # Recurse.
        return cls([cls.from_mask(m) for m in mask])

    def mask(self, dest=None):
        """Return the boolean mask equivalent of this object."""
        if dest is None:
            dest = np.empty(self.shape, bool)
        if len(self.ranges) and isinstance(self.ranges[0], Ranges):
            for d, r in zip(dest, self.ranges):
                d[:] = r.mask()
        else:
            # Recurse
            for d, rm in zip(dest, self.ranges):
                rm.mask(dest=d)
        return dest


# Support functions for RangesMatrix.__getitem__.  It's helpful to
# take this logic out of the class to handle some differences between
# Ranges and RangesMatrix.
#
# In _gibasic and _giadv, the target must be a Ranges or RangesMatrix,
# and index must be a tuple with no Ellipsis in it (but None are ok).
# The entry point is _gibasic, which will call _giadv if/when it
# encounteres an advanced index.

def _gibasic(target, index):
    if len(index) == 0:
        return target
    if index[0] is None:
        return RangesMatrix([_gibasic(target, index[1:])], skip_shape_check=True)
    is_rm = isinstance(target, RangesMatrix)
    if not is_rm and len(index) > 1:
        raise IndexError(f'Too many indices (extras: {index[1:]}).')
    if isinstance(index[0], (np.ndarray, list, tuple)):
        if is_rm:
            return _giadv(target, index)
        raise IndexError('Ranges (or last axis of RangesMatrix) '
                         'cannot use advanced indexing.')
    if isinstance(index[0], (int, np.int32, np.int64)):
        if is_rm:
            return _gibasic(target.ranges[index[0]], index[1:])
        raise IndexError('Cannot apply integer index to Ranges object.')
    if isinstance(index[0], slice):
        if is_rm:
            rm = RangesMatrix([_gibasic(r, index[1:]) for r in target.ranges[index[0]]],
                              child_shape=target.shape[1:], skip_shape_check=True)
            if rm.shape[0] == 0:
                # If your output doesn't have any .ranges, you need to
                # fake one in order to figure out how the dimensions
                # play out.
                fake_child = RangesMatrix.zeros(target.shape[1:])
                rm._child_shape = _gibasic(fake_child, index[1:]).shape
            return rm
        return target[index[0]]
    raise IndexError(f'Unexpected target[index]: {target}[{index}]')

def _giadv(target, index):
    adv_index, adv_axis, basic_index = [], [], []
    adr_axis = 0
    for axis, ind in enumerate(index):
        if isinstance(ind, (list, tuple, np.ndarray)):
            ind = np.asarray(ind)
            if ind.dtype == bool:
                if ind.ndim != 1 or len(ind) != target.shape[adr_axis]:
                    raise IndexError('index mask with shape '
                                     f'{ind.shape} is not compatible with '
                                     f'{target.shape[adr_axis]}')
                ind = ind.nonzero()[0]
            adv_index.append(ind)
            adv_axis.append(axis)
            basic_index.append(None)
        else:
            basic_index.append(ind)
        if ind is not None:
            adr_axis += 1
    assert(adv_axis[0] == 0)  # Don't call me until you need to.

    br = np.broadcast(*adv_index)

    # If zero size, short circuit.
    if br.size == 0:
        for _axis in adv_axis:
            basic_index[_axis] = 0
        child_thing = _gibasic(target, basic_index)
        return RangesMatrix.zeros(br.shape + child_thing.shape)

    # Note super_list will not be empty.
    super_list = []
    for idx_tuple in br:
        for _axis, _basic in zip(adv_axis, idx_tuple):
            basic_index[_axis] = _basic
        sub = _gibasic(target, basic_index)
        super_list.append(sub)

    return _giassemble(super_list, br.shape)

def _giassemble(items, shape):
    if len(shape) > 1:
        stride = len(items) // shape[0]
        groups = [items[i*stride:(i+1)*stride] for i in range(shape[0])]
        items = [_giassemble(g, shape[1:]) for g in groups]
    return RangesMatrix(items)

1	import so3g	1✔
2	import numpy as np	1✔
3
4	"""Objects will self report as being of type "RangesInt32" rather than	1✔
5	Ranges. But let's try to use so3g.proj.Ranges when testing types and
6	making new ones and stuff."""
7
8	Ranges = so3g.RangesInt32	1✔
9
10
11	class RangesMatrix():	1✔
12	"""This is a wrapper for multi-dimensional grids of Ranges objects.
13	This can be used to store Ranges objects per-detector (the 2-d
14	case) or per-thread and per-detector (a 3-d case, required if
15	using OpenMP). The right-most axis always corresponds to the
16	sample axis carried by the Ranges objects.
17
18	In addition to .shape and multi-dimensional slicing, it supports
19	inversion (the ~ operator) and multiplication against other
20	RangesMatrix or Ranges objects, with broadcasting rules similar to
21	standard numpy array rules.
22
23	"""
24	def __init__(self, items=[], child_shape=None, skip_shape_check=False):	1✔
25	self.ranges = [x for x in items]	1✔
26	if len(items):	1✔
27	child_shape = items[0].shape	1✔
28	if not skip_shape_check:	1✔
29	assert all([(item.shape == child_shape) for item in items])	1✔
30	elif child_shape is None:	1✔
31	child_shape = ()	×
32	self._child_shape = child_shape	1✔
33
34	def __repr__(self):	1✔
35	return 'RangesMatrix(' + ','.join(map(str, self.shape)) + ')'	1✔
36
37	def __len__(self):	1✔
38	return self.shape[0]	1✔
39
40	def copy(self):	1✔
41	return RangesMatrix([x.copy() for x in self.ranges],	1✔
42	child_shape=self.shape[1:])
43
44	def zeros_like(self):	1✔
45	return RangesMatrix([x.zeros_like() for x in self.ranges],	×
46	child_shape=self.shape[1:])
47
48	def ones_like(self):	1✔
49	return RangesMatrix([x.ones_like() for x in self.ranges],	×
50	child_shape=self.shape[1:])
51
52	def buffer(self, buff):	1✔
53	[x.buffer(buff) for x in self.ranges]	×
54	## just to make this work like Ranges.buffer()
55	return self	×
56
57	def buffered(self, buff):	1✔
58	out = self.copy()	×
59	[x.buffer(buff) for x in out.ranges]	×
60	return out	×
61
62	@property	1✔
63	def shape(self):	1✔
64	if len(self.ranges) == 0:	1✔
65	return (0,) + self._child_shape	1✔
66	return (len(self.ranges),) + self.ranges[0].shape	1✔
67
68	def __getitem__(self, index):	1✔
69	"""RangesMatrix supports multi-dimensional indexing, slicing, and
70	numpy-style advanced indexing (with some restrictions).
71
72	The right-most axis of RangesMatrix has special restrictions,
73	since it corresponds to Ranges objects: it can only accept
74	slice-based indexing, not integer or advanced indexing.
75
76	To guarantee that you get copies, rather than references to,
77	the lowest level Ranges objects, make sure the index tuple
78	includes a slice along the final (right-most) axis. For
79	example::
80
81	# Starting point
82	rm = RangesMatrix.zeros(10, 10000)
83
84	# Get and modify an element in rm ...
85	r = rm[0]
86	r.add_interval(0, 10) # This _does_ affect rm!
87
88	# Get a copy of an element from rm.
89	r = rm[0,:]
90	r.add_interval(0, 10) # This will not affect rm.
91
92	# More generally, this is equivalent to rm1 = rm.copy():
93	new_rm = rm[..., :]
94
95	"""
96	if not isinstance(index, tuple):	1✔
97	index = (index,)	1✔
98
99	eidx = [i for i, a in enumerate(index) if a is Ellipsis]	1✔
100	if len(eidx) == 1:	1✔
101	# Fill in missing slices.
102	eidx = eidx[0]	1✔
103	n_free = len(self.shape) - sum([e is not None for e in index]) + 1	1✔
104	index = index[:eidx] + tuple([slice(None)] * n_free) + index[eidx+1:]	1✔
105	elif len(eidx) > 1:	1✔
106	raise IndexError("An index can only have a single ellipsis ('...')")	×
107
108	return _gibasic(self, index)	1✔
109
110	def __add__(self, x):	1✔
111	if isinstance(x, Ranges):	×
NEW 112	return self.__class__([d + x for d in self.ranges], skip_shape_check=True)	×
113	elif isinstance(x, RangesMatrix):	×
114	# Check for shape compatibility.
115	nd_a = len(self.shape)	×
116	nd_b = len(x.shape)	×
117	ndim = min(nd_a, nd_b)	×
118	ok = [(a==1 or b==1 or a == b) for a,b in zip(self.shape[-ndim:], x.shape[-ndim:])]	×
119	if not all(ok) or nd_a < nd_b:	×
120	raise ValueError('Operands have incompatible shapes: %s %s' %	×
121	self.shape, x.shape)
122	if nd_a == nd_b:	×
123	# Broadcast if either has shape 1...
124	if x.shape[0] == 1:	×
NEW 125	return self.__class__([r + x[0] for r in self.ranges], skip_shape_check=True)	×
126	elif self.shape[0] == 1:	×
NEW 127	return self.__class__([self.ranges[0] + _x for _x in x], skip_shape_check=True)	×
128	elif self.shape[0] == x.shape[0]:	×
NEW 129	return self.__class__([r + d for r, d in zip(self.ranges, x)], skip_shape_check=True)	×
NEW 130	return self.__class__([r + x for r in self.ranges], skip_shape_check=True)	×
131
132	def __mul__(self, x):	1✔
133	if isinstance(x, Ranges):	×
NEW 134	return self.__class__([d * x for d in self.ranges], skip_shape_check=True)	×
135	elif isinstance(x, RangesMatrix):	×
136	# Check for shape compatibility.
137	nd_a = len(self.shape)	×
138	nd_b = len(x.shape)	×
139	ndim = min(nd_a, nd_b)	×
140	ok = [(a==1 or b==1 or a == b) for a,b in zip(self.shape[-ndim:], x.shape[-ndim:])]	×
141	if not all(ok) or nd_a < nd_b:	×
142	raise ValueError('Operands have incompatible shapes: %s %s' %	×
143	self.shape, x.shape)
144	if nd_a == nd_b:	×
145	# Broadcast if either has shape 1...
146	if x.shape[0] == 1:	×
NEW 147	return self.__class__([r * x[0] for r in self.ranges], skip_shape_check=True)	×
148	elif self.shape[0] == 1:	×
NEW 149	return self.__class__([self.ranges[0] * _x for _x in x], skip_shape_check=True)	×
150	elif self.shape[0] == x.shape[0]:	×
NEW 151	return self.__class__([r * d for r, d in zip(self.ranges, x)], skip_shape_check=True)	×
NEW 152	return self.__class__([r * x for r in self.ranges], skip_shape_check=True)	×
153
154	def __invert__(self):	1✔
155	return self.__class__([~x for x in self.ranges], skip_shape_check=True)	1✔
156
157	@staticmethod	1✔
158	def concatenate(items, axis=0):	1✔
159	"""Static method to concatenate multiple RangesMatrix (or Ranges)
160	objects along the specified axis.
161
162	"""
163	# Check shape compatibility...
164	s = list(items[0].shape)	1✔
165	s[axis] = -1	1✔
166	for item in items[1:]:	1✔
167	s1 = list(item.shape)	1✔
168	s1[axis] = -1	1✔
169	if s != s1:	1✔
170	raise ValueError('Contributed items must have same shape on non-cat axis.')	1✔
171	def collect(items, join_depth):	1✔
172	# Recurse through items in order to concatenate along axis join_depth.
173	ranges = []	1✔
174	if join_depth > 0:	1✔
175	for co_items in zip(*items):	1✔
176	ranges.append(collect(co_items, join_depth - 1))	1✔
177	else:
178	if isinstance(items[0], RangesMatrix):	1✔
179	for item in items:	1✔
180	ranges.extend(item.ranges)	1✔
181	else:
182	# The items are Ranges, then.
183	n, r = 0, []	1✔
184	for item in items:	1✔
185	r.extend(item.ranges() + n)	1✔
186	n += item.count	1✔
187	r = Ranges.from_array(	1✔
188	np.array(r, dtype='int32').reshape((-1, 2)), n)
189	return r	1✔
190	return RangesMatrix(ranges, child_shape=items[0].shape[1:])	1✔
191	return collect(items, axis)	1✔
192
193	def close_gaps(self, gap_size=0):	1✔
194	"""Call close_gaps(gap_size) on all children. Any ranges that abutt
195	each other within the gap_size are merged into a single entry.
196	Usually a gap_size of 0 is not possible, but if a caller is
197	carefully using append_interval_no_check, then it can happen.
198
199	"""
200	for r in self.ranges:	1✔
201	r.close_gaps(gap_size)	1✔
202
203	def get_stats(self):	1✔
204	samples, intervals = [], []	×
205	for r in self.ranges:	×
206	if isinstance(r, Ranges):	×
207	ra = r.ranges()	×
208	samples.append(np.dot(ra, [-1, 1]).sum())	×
209	intervals.append(ra.shape[0])	×
210	else:
211	sub = r.get_stats()	×
212	samples.append(sum(sub['samples']))	×
213	intervals.append(sum(sub['intervals']))	×
214	return {	×
215	'samples': samples,
216	'intervals': intervals}
217
218	@staticmethod	1✔
219	def full(shape, fill_value):	1✔
220	"""Construct a RangesMatrix with the specified shape, initialized to
221	fill_value.
222
223	Args:
224	shape (tuple of int): The shape. Must have at least one
225	element. If there is only one element, a Ranges object is
226	returned, not a RangesMatrix.
227	fill_value (bool): True or False. If False, the wrapped
228	Ranges objects will have no intervals. If True, the
229	wrapped Ranges objects will all have a single interval
230	spanning their entire domain.
231
232	Returns:
233	Ranges object (if shape has a single element) or a
234	RangesMatrix object (len(shape) > 1).
235
236	See also: zeros, ones.
237
238	"""
239	assert fill_value in [True, False]	1✔
240	if isinstance(shape, int):	1✔
241	shape = (shape,)	×
242	assert(len(shape) > 0)	1✔
243	if len(shape) == 1:	1✔
244	r = Ranges(shape[0])	1✔
245	if fill_value:	1✔
246	r = ~r	1✔
247	return r	1✔
248	return RangesMatrix([RangesMatrix.full(shape[1:], fill_value)	1✔
249	for i in range(shape[0])],
250	child_shape=shape[1:])
251
252	@classmethod	1✔
253	def zeros(cls, shape):	1✔
254	"""Equivalent to full(shape, False)."""
255	return cls.full(shape, False)	1✔
256
257	@classmethod	1✔
258	def ones(cls, shape):	1✔
259	"""Equivalent to full(shape, True)."""
260	return cls.full(shape, True)	1✔
261
262	@classmethod	1✔
263	def from_mask(cls, mask):	1✔
264	"""Create a RangesMatrix from a boolean mask. The input mask can have
265	any dimensionality greater than 0 but be aware that if ndim==1
266	then a Ranges object, and not a RangesMatrix, is returned.
267
268	Args:
269	mask (ndarray): Must be boolean array with at least 1 dimension.
270
271	Returns:
272	RangesMatrix with the same shape as mask, with ranges
273	corresponding to the intervals where mask was True.
274
275	"""
276	assert(mask.ndim > 0)	1✔
277	if mask.ndim == 1:	1✔
278	return Ranges.from_mask(mask)	1✔
279	if len(mask) == 0:	1✔
280	return cls(child_shape=mask.shape[1:])	1✔
281	# Recurse.
282	return cls([cls.from_mask(m) for m in mask])	1✔
283
284	def mask(self, dest=None):	1✔
285	"""Return the boolean mask equivalent of this object."""
286	if dest is None:	1✔
287	dest = np.empty(self.shape, bool)	1✔
288	if len(self.ranges) and isinstance(self.ranges[0], Ranges):	1✔
289	for d, r in zip(dest, self.ranges):	1✔
290	d[:] = r.mask()	1✔
291	else:
292	# Recurse
293	for d, rm in zip(dest, self.ranges):	1✔
294	rm.mask(dest=d)	1✔
295	return dest	1✔
296
297
298	# Support functions for RangesMatrix.__getitem__. It's helpful to
299	# take this logic out of the class to handle some differences between
300	# Ranges and RangesMatrix.
301	#
302	# In _gibasic and _giadv, the target must be a Ranges or RangesMatrix,
303	# and index must be a tuple with no Ellipsis in it (but None are ok).
304	# The entry point is _gibasic, which will call _giadv if/when it
305	# encounteres an advanced index.
306
307	def _gibasic(target, index):	1✔
308	if len(index) == 0:	1✔
309	return target	1✔
310	if index[0] is None:	1✔
311	return RangesMatrix([_gibasic(target, index[1:])], skip_shape_check=True)	1✔
312	is_rm = isinstance(target, RangesMatrix)	1✔
313	if not is_rm and len(index) > 1:	1✔
314	raise IndexError(f'Too many indices (extras: {index[1:]}).')	1✔
315	if isinstance(index[0], (np.ndarray, list, tuple)):	1✔
316	if is_rm:	1✔
317	return _giadv(target, index)	1✔
318	raise IndexError('Ranges (or last axis of RangesMatrix) '	1✔
319	'cannot use advanced indexing.')
320	if isinstance(index[0], (int, np.int32, np.int64)):	1✔
321	if is_rm:	1✔
322	return _gibasic(target.ranges[index[0]], index[1:])	1✔
323	raise IndexError('Cannot apply integer index to Ranges object.')	1✔
324	if isinstance(index[0], slice):	1✔
325	if is_rm:	1✔
326	rm = RangesMatrix([_gibasic(r, index[1:]) for r in target.ranges[index[0]]],	1✔
327	child_shape=target.shape[1:], skip_shape_check=True)
328	if rm.shape[0] == 0:	1✔
329	# If your output doesn't have any .ranges, you need to
330	# fake one in order to figure out how the dimensions
331	# play out.
332	fake_child = RangesMatrix.zeros(target.shape[1:])	1✔
333	rm._child_shape = _gibasic(fake_child, index[1:]).shape	1✔
334	return rm	1✔
335	return target[index[0]]	1✔
336	raise IndexError(f'Unexpected target[index]: {target}[{index}]')	×
337
338	def _giadv(target, index):	1✔
339	adv_index, adv_axis, basic_index = [], [], []	1✔
340	adr_axis = 0	1✔
341	for axis, ind in enumerate(index):	1✔
342	if isinstance(ind, (list, tuple, np.ndarray)):	1✔
343	ind = np.asarray(ind)	1✔
344	if ind.dtype == bool:	1✔
345	if ind.ndim != 1 or len(ind) != target.shape[adr_axis]:	1✔
346	raise IndexError('index mask with shape '	1✔
347	f'{ind.shape} is not compatible with '
348	f'{target.shape[adr_axis]}')
349	ind = ind.nonzero()[0]	1✔
350	adv_index.append(ind)	1✔
351	adv_axis.append(axis)	1✔
352	basic_index.append(None)	1✔
353	else:
354	basic_index.append(ind)	1✔
355	if ind is not None:	1✔
356	adr_axis += 1	1✔
357	assert(adv_axis[0] == 0) # Don't call me until you need to.	1✔
358
359	br = np.broadcast(*adv_index)	1✔
360
361	# If zero size, short circuit.
362	if br.size == 0:	1✔
363	for _axis in adv_axis:	1✔
364	basic_index[_axis] = 0	1✔
365	child_thing = _gibasic(target, basic_index)	1✔
366	return RangesMatrix.zeros(br.shape + child_thing.shape)	1✔
367
368	# Note super_list will not be empty.
369	super_list = []	1✔
370	for idx_tuple in br:	1✔
371	for _axis, _basic in zip(adv_axis, idx_tuple):	1✔
372	basic_index[_axis] = _basic	1✔
373	sub = _gibasic(target, basic_index)	1✔
374	super_list.append(sub)	1✔
375
376	return _giassemble(super_list, br.shape)	1✔
377
378	def _giassemble(items, shape):	1✔
379	if len(shape) > 1:	1✔
380	stride = len(items) // shape[0]	1✔
381	groups = [items[istride:(i+1)stride] for i in range(shape[0])]	1✔
382	items = [_giassemble(g, shape[1:]) for g in groups]	1✔
383	return RangesMatrix(items)	1✔

simonsobs / so3g / 22156500136

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous