sgnl.transforms.lloid_correlate

An element that performs cross-correlation of input data with filters, with certain time delays and padding for upsampling and downsampling

LLOIDCorrelate dataclass

Bases: TSTransform

Correlates input data with filters, with certain time delays, and padding for upsampling and downsampling

Parameters:

Name	Type	Description	Default
filters	Array	Array, the filter to correlate over	None
uppad	int	int, the upsampling padding, in offsets	0
downpad	int	int, the downsampling padding, in offsets	0
delays	list[int]	list[int], a list of time delays for each time slice, in offsets	None

Source code in sgnl/transforms/lloid_correlate.py

@dataclass
class LLOIDCorrelate(TSTransform):
    """Correlates input data with filters, with certain time delays, and padding for
    upsampling and downsampling

    Args:
        filters:
            Array, the filter to correlate over
        uppad:
            int, the upsampling padding, in offsets
        downpad:
            int, the downsampling padding, in offsets
        delays:
            list[int], a list of time delays for each time slice, in offsets
    """

    # Torch-only: the correlation is a torch ``conv1d`` over torch filters.
    backends = frozenset({"torch"})

    # FFT/conv1d dispatch threshold (plain class attribute, not a dataclass
    # field). Direct valid correlation costs ~L*tout MACs per filter while the
    # FFT path costs ~nfft*log2(nfft) per filter regardless of tout, so FFT is
    # used when L*tout >= factor * nfft * log2(nfft). Calibrated on CPU (torch
    # 2.5): measured break-even is a ratio of ~20 across filter/block lengths
    # (conv1d wins 2x at ratio ~11, FFT wins 2.3x at ratio ~43).
    _FFT_COST_FACTOR = 20

    filters: Array = None
    uppad: int = 0
    downpad: int = 0
    delays: list[int] = None  # type: ignore[assignment]
    reconstruction_segment_list: segments.segmentlist | None = None

    def __post_init__(self):
        assert self.filters is not None
        assert self.delays is not None, "delays must be provided"
        self.shape = self.filters.shape
        self.filters = self.filters.view(-1, 1, self.shape[-1])
        super().__post_init__()
        assert (
            len(self.sink_pads) == 1 and len(self.source_pads) == 1
        ), "only one sink_pad and one source_pad is allowed"

        self._fft_cache = None
        self.audioadapter = Audioadapter()
        self.unique_delays = sorted(set(self.delays))
        self.startup = True

    def corr(self, data: Array) -> Array:
        """Correlate the data with filters.

        Dispatches per block shape between FFT correlation with cached filter
        spectra and grouped time-domain conv1d, whichever is cheaper: the FFT
        transforms cost ~nfft*log2(nfft) per filter regardless of the output
        length, so they win for the long blocks of the max-rate LLOID groups
        (~8x on CPU) but lose to direct valid correlation (~L*tout MACs per
        filter) for the short blocks of the low-rate groups. The FFT path is
        exact on the valid range (no circular contamination because
        nfft - L + 1 >= tout); the decision and the filter spectra are cached
        per (channels, block-length).

        Args:
            data:
                Array, the data to correlate with the filters

        Returns:
            Array, the output of the correlation
        """
        L = self.filters.shape[-1]
        C, T = data.shape[-2], data.shape[-1]
        tout = T - L + 1
        if data.dim() not in (2, 3) or tout < 1:
            # FIXME: try supporting numpy?
            return Fconv1d(data, self.filters, groups=data.shape[-2]).view(
                self.shape[:-1] + (-1,)
            )
        if self._fft_cache is None or self._fft_cache[0] != (C, T):
            nfft = 1 << (T - 1).bit_length()
            if L * tout >= self._FFT_COST_FACTOR * nfft * (nfft.bit_length() - 1):
                wf = torch.fft.rfft(self.filters.reshape(C, -1, L).flip(-1), n=nfft)
            else:
                # Output block too short for the transforms to pay for
                # themselves; wf None marks the conv1d choice in the cache.
                wf = None
            self._fft_cache = ((C, T), nfft, wf)
        _, nfft, wf = self._fft_cache
        if wf is None:
            return Fconv1d(data, self.filters, groups=data.shape[-2]).view(
                self.shape[:-1] + (-1,)
            )
        xf = torch.fft.rfft(data, n=nfft)
        prod = xf.unsqueeze(-2) * wf
        out = torch.fft.irfft(prod, n=nfft)[..., L - 1 : L - 1 + tout]
        return out.reshape(self.shape[:-1] + (-1,))

    def new(self, pad: SourcePad) -> TSFrame:
        """Correlates incoming frames with filters.

        Args:
            pad:
                SourcePad, the source pad to produce the transformed frames

        Returns:
            TSFrame, the output TSFrame
        """
        A = self.audioadapter
        frame = self.preparedframes[self.sink_pads[0]]
        assert isinstance(frame, TSFrame)
        outbufs = []

        # process buffer by buffer
        for buf in frame:
            # find the reference segment "this_segment"
            if buf.end_offset - buf.offset == 0:
                if self.startup:
                    this_segment1 = buf.end_offset
                else:
                    this_segment1 = buf.end_offset + self.downpad
                this_segment0 = this_segment1
                outbufs.append(
                    SeriesBuffer(
                        offset=this_segment0 + self.uppad,
                        sample_rate=buf.sample_rate,
                        data=None,
                        shape=self.shape[:-1] + (0,),
                    )
                )
            else:
                this_segment1 = buf.end_offset + self.downpad
                if self.startup:
                    this_segment0 = (
                        this_segment1 - (buf.end_offset - buf.offset) - self.downpad
                    )
                    self.startup = False
                else:
                    this_segment0 = this_segment1 - (buf.end_offset - buf.offset)
                this_segment = (this_segment0, this_segment1)
                A.push(buf)

                if (
                    self.reconstruction_segment_list is not None
                    and not self.reconstruction_segment_list.intersects_segment(
                        segments.segment(
                            LIGOTimeGPS(Offset.tosec(this_segment0)),
                            LIGOTimeGPS(Offset.tosec(this_segment1)),
                        )
                    )
                ):
                    # The ortho snr segment does not intersect the reconstruction
                    # segment list, don't need to produce this ortho snr
                    outbufs.append(
                        SeriesBuffer(
                            offset=this_segment[0] + self.uppad,
                            sample_rate=buf.sample_rate,
                            data=None,
                            shape=self.shape[:-1]
                            + (
                                Offset.tosamples(
                                    this_segment1 - this_segment0, buf.sample_rate
                                ),
                            ),
                        )
                    )
                else:
                    outs_map = {}
                    # Only do the copy for unique delays
                    copied_data = False
                    earliest = []

                    # copy out the unique segments
                    for delay in self.unique_delays:
                        # find the segment to copy out
                        cp_segment1 = this_segment1 + self.uppad - delay
                        cp_segment0 = (
                            cp_segment1
                            - (this_segment1 - this_segment0)
                            - Offset.fromsamples(self.shape[-1] - 1, buf.sample_rate)
                        )
                        earliest.append(cp_segment0)
                        if cp_segment1 > A.offset and not A.is_gap:
                            if A.pre_cat_data is None:
                                A.concatenate_data(
                                    (
                                        max(
                                            this_segment0
                                            + self.uppad
                                            - max(self.unique_delays)
                                            - Offset.fromsamples(
                                                self.shape[-1] - 1, buf.sample_rate
                                            ),
                                            A.offset,
                                        ),
                                        this_segment1
                                        + self.uppad
                                        - min(self.unique_delays),
                                    )
                                )
                            cp_segment = (max(A.offset, cp_segment0), cp_segment1)
                            # We need to do a copy
                            out = A.copy_samples_by_offset_segment(cp_segment)
                            if cp_segment0 < A.offset and out is not None:
                                # pad with zeros in front
                                pad_length = Offset.tosamples(
                                    A.offset - cp_segment0, buf.sample_rate
                                )
                                out = signal_pad(out, (pad_length, 0))
                            # the copy returns None when the requested segment
                            # lies entirely within gap buffers, even though the
                            # adapter as a whole is not gap
                            if out is not None:
                                copied_data = True
                        else:
                            out = None
                        outs_map[delay] = out

                    # fill in zeros arrays
                    if copied_data is True:
                        # Now stack the output array
                        if len(self.unique_delays) == 1:
                            # copied_data is True, so the single delay's copy
                            # is real data
                            outs: Array | None = outs_map[
                                self.unique_delays[0]
                            ].unsqueeze(0)
                        else:
                            # The gaps must match the copied data's backend; use a
                            # real copied array as the reference ("the array is the
                            # backend"). copied_data is True, so one exists.
                            reference = next(
                                o for o in outs_map.values() if o is not None
                            )
                            xp = array_namespace(reference)
                            assert xp is not None  # reference is a real array
                            outs_list: list[Array] = []
                            for delay in self.delays:
                                out = outs_map[delay]
                                if out is None:
                                    out = new_zeros(
                                        reference,
                                        (
                                            Offset.tosamples(
                                                cp_segment1 - cp_segment0,
                                                buf.sample_rate,
                                            ),
                                        ),
                                    )
                                outs_list.append(out)
                            outs = xp.stack(outs_list)
                    else:
                        outs = None

                    # flush data
                    flush_end_offset = min(earliest)
                    if flush_end_offset > A.offset:
                        A.flush_samples_by_end_offset(flush_end_offset)

                    # Do the correlation!
                    if outs is not None:
                        outs = self.corr(outs)
                    outbufs.append(
                        SeriesBuffer(
                            offset=this_segment[0] + self.uppad,
                            sample_rate=buf.sample_rate,
                            data=outs,
                            shape=self.shape[:-1]
                            + (
                                Offset.tosamples(
                                    this_segment1 - this_segment0, buf.sample_rate
                                ),
                            ),
                        )
                    )

        # This element builds its output frame by hand, so stamp the upstream
        # backend identity: an all-gap output keeps the (torch) stream, not numpy.
        out_frame = TSFrame(
            buffers=outbufs,
            EOS=frame.EOS,
            metadata=frame.metadata,
            spec_hint=self._output_hint(pad),
        )

        return out_frame

corr(data)

Correlate the data with filters.

Dispatches per block shape between FFT correlation with cached filter spectra and grouped time-domain conv1d, whichever is cheaper: the FFT transforms cost ~nfftlog2(nfft) per filter regardless of the output length, so they win for the long blocks of the max-rate LLOID groups (~8x on CPU) but lose to direct valid correlation (~Ltout MACs per filter) for the short blocks of the low-rate groups. The FFT path is exact on the valid range (no circular contamination because nfft - L + 1 >= tout); the decision and the filter spectra are cached per (channels, block-length).

Parameters:

Name	Type	Description	Default
data	Array	Array, the data to correlate with the filters	required

Returns:

Type	Description
Array	Array, the output of the correlation

Source code in sgnl/transforms/lloid_correlate.py

def corr(self, data: Array) -> Array:
    """Correlate the data with filters.

    Dispatches per block shape between FFT correlation with cached filter
    spectra and grouped time-domain conv1d, whichever is cheaper: the FFT
    transforms cost ~nfft*log2(nfft) per filter regardless of the output
    length, so they win for the long blocks of the max-rate LLOID groups
    (~8x on CPU) but lose to direct valid correlation (~L*tout MACs per
    filter) for the short blocks of the low-rate groups. The FFT path is
    exact on the valid range (no circular contamination because
    nfft - L + 1 >= tout); the decision and the filter spectra are cached
    per (channels, block-length).

    Args:
        data:
            Array, the data to correlate with the filters

    Returns:
        Array, the output of the correlation
    """
    L = self.filters.shape[-1]
    C, T = data.shape[-2], data.shape[-1]
    tout = T - L + 1
    if data.dim() not in (2, 3) or tout < 1:
        # FIXME: try supporting numpy?
        return Fconv1d(data, self.filters, groups=data.shape[-2]).view(
            self.shape[:-1] + (-1,)
        )
    if self._fft_cache is None or self._fft_cache[0] != (C, T):
        nfft = 1 << (T - 1).bit_length()
        if L * tout >= self._FFT_COST_FACTOR * nfft * (nfft.bit_length() - 1):
            wf = torch.fft.rfft(self.filters.reshape(C, -1, L).flip(-1), n=nfft)
        else:
            # Output block too short for the transforms to pay for
            # themselves; wf None marks the conv1d choice in the cache.
            wf = None
        self._fft_cache = ((C, T), nfft, wf)
    _, nfft, wf = self._fft_cache
    if wf is None:
        return Fconv1d(data, self.filters, groups=data.shape[-2]).view(
            self.shape[:-1] + (-1,)
        )
    xf = torch.fft.rfft(data, n=nfft)
    prod = xf.unsqueeze(-2) * wf
    out = torch.fft.irfft(prod, n=nfft)[..., L - 1 : L - 1 + tout]
    return out.reshape(self.shape[:-1] + (-1,))

new(pad)

Correlates incoming frames with filters.

Parameters:

Name	Type	Description	Default
pad	SourcePad	SourcePad, the source pad to produce the transformed frames	required

Returns:

Type	Description
TSFrame	TSFrame, the output TSFrame

Source code in sgnl/transforms/lloid_correlate.py

def new(self, pad: SourcePad) -> TSFrame:
    """Correlates incoming frames with filters.

    Args:
        pad:
            SourcePad, the source pad to produce the transformed frames

    Returns:
        TSFrame, the output TSFrame
    """
    A = self.audioadapter
    frame = self.preparedframes[self.sink_pads[0]]
    assert isinstance(frame, TSFrame)
    outbufs = []

    # process buffer by buffer
    for buf in frame:
        # find the reference segment "this_segment"
        if buf.end_offset - buf.offset == 0:
            if self.startup:
                this_segment1 = buf.end_offset
            else:
                this_segment1 = buf.end_offset + self.downpad
            this_segment0 = this_segment1
            outbufs.append(
                SeriesBuffer(
                    offset=this_segment0 + self.uppad,
                    sample_rate=buf.sample_rate,
                    data=None,
                    shape=self.shape[:-1] + (0,),
                )
            )
        else:
            this_segment1 = buf.end_offset + self.downpad
            if self.startup:
                this_segment0 = (
                    this_segment1 - (buf.end_offset - buf.offset) - self.downpad
                )
                self.startup = False
            else:
                this_segment0 = this_segment1 - (buf.end_offset - buf.offset)
            this_segment = (this_segment0, this_segment1)
            A.push(buf)

            if (
                self.reconstruction_segment_list is not None
                and not self.reconstruction_segment_list.intersects_segment(
                    segments.segment(
                        LIGOTimeGPS(Offset.tosec(this_segment0)),
                        LIGOTimeGPS(Offset.tosec(this_segment1)),
                    )
                )
            ):
                # The ortho snr segment does not intersect the reconstruction
                # segment list, don't need to produce this ortho snr
                outbufs.append(
                    SeriesBuffer(
                        offset=this_segment[0] + self.uppad,
                        sample_rate=buf.sample_rate,
                        data=None,
                        shape=self.shape[:-1]
                        + (
                            Offset.tosamples(
                                this_segment1 - this_segment0, buf.sample_rate
                            ),
                        ),
                    )
                )
            else:
                outs_map = {}
                # Only do the copy for unique delays
                copied_data = False
                earliest = []

                # copy out the unique segments
                for delay in self.unique_delays:
                    # find the segment to copy out
                    cp_segment1 = this_segment1 + self.uppad - delay
                    cp_segment0 = (
                        cp_segment1
                        - (this_segment1 - this_segment0)
                        - Offset.fromsamples(self.shape[-1] - 1, buf.sample_rate)
                    )
                    earliest.append(cp_segment0)
                    if cp_segment1 > A.offset and not A.is_gap:
                        if A.pre_cat_data is None:
                            A.concatenate_data(
                                (
                                    max(
                                        this_segment0
                                        + self.uppad
                                        - max(self.unique_delays)
                                        - Offset.fromsamples(
                                            self.shape[-1] - 1, buf.sample_rate
                                        ),
                                        A.offset,
                                    ),
                                    this_segment1
                                    + self.uppad
                                    - min(self.unique_delays),
                                )
                            )
                        cp_segment = (max(A.offset, cp_segment0), cp_segment1)
                        # We need to do a copy
                        out = A.copy_samples_by_offset_segment(cp_segment)
                        if cp_segment0 < A.offset and out is not None:
                            # pad with zeros in front
                            pad_length = Offset.tosamples(
                                A.offset - cp_segment0, buf.sample_rate
                            )
                            out = signal_pad(out, (pad_length, 0))
                        # the copy returns None when the requested segment
                        # lies entirely within gap buffers, even though the
                        # adapter as a whole is not gap
                        if out is not None:
                            copied_data = True
                    else:
                        out = None
                    outs_map[delay] = out

                # fill in zeros arrays
                if copied_data is True:
                    # Now stack the output array
                    if len(self.unique_delays) == 1:
                        # copied_data is True, so the single delay's copy
                        # is real data
                        outs: Array | None = outs_map[
                            self.unique_delays[0]
                        ].unsqueeze(0)
                    else:
                        # The gaps must match the copied data's backend; use a
                        # real copied array as the reference ("the array is the
                        # backend"). copied_data is True, so one exists.
                        reference = next(
                            o for o in outs_map.values() if o is not None
                        )
                        xp = array_namespace(reference)
                        assert xp is not None  # reference is a real array
                        outs_list: list[Array] = []
                        for delay in self.delays:
                            out = outs_map[delay]
                            if out is None:
                                out = new_zeros(
                                    reference,
                                    (
                                        Offset.tosamples(
                                            cp_segment1 - cp_segment0,
                                            buf.sample_rate,
                                        ),
                                    ),
                                )
                            outs_list.append(out)
                        outs = xp.stack(outs_list)
                else:
                    outs = None

                # flush data
                flush_end_offset = min(earliest)
                if flush_end_offset > A.offset:
                    A.flush_samples_by_end_offset(flush_end_offset)

                # Do the correlation!
                if outs is not None:
                    outs = self.corr(outs)
                outbufs.append(
                    SeriesBuffer(
                        offset=this_segment[0] + self.uppad,
                        sample_rate=buf.sample_rate,
                        data=outs,
                        shape=self.shape[:-1]
                        + (
                            Offset.tosamples(
                                this_segment1 - this_segment0, buf.sample_rate
                            ),
                        ),
                    )
                )

    # This element builds its output frame by hand, so stamp the upstream
    # backend identity: an all-gap output keeps the (torch) stream, not numpy.
    out_frame = TSFrame(
        buffers=outbufs,
        EOS=frame.EOS,
        metadata=frame.metadata,
        spec_hint=self._output_hint(pad),
    )

    return out_frame