Skip to content

sgnts.transforms.sumindex

SumIndex dataclass

Bases: TSTransform

Sum array values over slices in the zero-th dimension.

Parameters:

Name Type Description Default
sl Optional[list[slice]]

list[slice], the slices to sum over

 None
backend type[ArrayBackend]

type[ArrayBackend], the wrapper around array operations.

 NumpyBackend
Source code in sgnts/transforms/sumindex.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
@dataclass
class SumIndex(TSTransform):
    """Sum array values over slices in the zero-th dimension.

    Args:
        sl:
            list[slice], the slices to sum over
        backend:
            type[ArrayBackend], the wrapper around array operations.
    """

    sl: Optional[list[slice]] = None
    backend: type[ArrayBackend] = NumpyBackend

    def __post_init__(self):
        super().__post_init__()
        assert self.sl is not None
        for sl in self.sl:
            assert isinstance(sl, slice)

    # FIXME: wraps are not playing well with mypy.  For now ignore and hope
    # that a future version of mypy will be able to handle this
    @wraps(TSTransform.new)
    def new(self, pad: SourcePad) -> TSFrame:  # type: ignore
        frame = self.preparedframes[self.sink_pads[0]]

        outbufs = []
        for buf in frame:
            if buf.is_gap:
                out = None
            else:
                data = buf.data
                data_all = []
                # NOTE mypy complains about None not being iterable but None
                # should actually be impossible at this point.
                assert self.sl is not None
                for sl in self.sl:
                    if sl.stop - sl.start == 1:
                        data_all.append((data[sl.start, :, :]))
                    else:
                        data_all.append(self.backend.sum(data[sl, :, :], axis=0))

                out = self.backend.stack(data_all)

            # NOTE mypy complains about None not being iterable but None should
            # actually be impossible at this point.
            assert self.sl is not None
            outbuf = SeriesBuffer(
                offset=buf.offset,
                sample_rate=buf.sample_rate,
                data=out,
                shape=(len(self.sl),) + buf.shape[-2:],
            )
        outbufs.append(outbuf)

        return TSFrame(buffers=outbufs, EOS=frame.EOS, metadata=frame.metadata)