ezmsg-org / ezmsg

Pure-Python DAG-based high-performance SHM-backed pub-sub and multi-processing pattern
https://ezmsg.readthedocs.io/en/latest/
MIT License
9 stars 5 forks source link

bug(?): sigproc.window, if newaxis, should still update offset in old axis #58

Closed cboulay closed 5 months ago

cboulay commented 7 months ago

https://github.com/iscoe/ezmsg/blob/43dda8bf724d39b4cfa6094187bf59b70483ef66/extensions/ezmsg-sigproc/ezmsg/sigproc/window.py#L121-L144

When using sigproc.Window, I use the newaxis="step" setting to make sure I can keep track of the time in axes["step"].offset, even after a subsequent Spectrum operation converts the original 'time' axis to 'freq'.

However, when I do this, the original time axis is no longer modified to reflect the windowing.

I think that the linked code should be replaced with the following:

        for out_view, offset in outputs:
            out_view = np.moveaxis(out_view, 0, axis_idx)
            out_dims = msg.dims
            out_axes = msg.axes
            if axis_name in msg.axes:    
                out_axes[axis_name] = replace(axis, offset=offset)

            if (
                self.STATE.cur_settings.newaxis is not None
                and self.STATE.cur_settings.newaxis != self.STATE.cur_settings.axis
            ):
                out_view = out_view[np.newaxis, ...]
                out_dims = [self.STATE.cur_settings.newaxis] + msg.dims

                new_gain = 0.0
                if self.STATE.window_shift_samples is not None:
                    new_gain = axis.gain * self.STATE.window_shift_samples
                new_axis = AxisArray.Axis(unit=axis.unit, gain=new_gain, offset=offset)
                out_axes = {**msg.axes, **{self.STATE.cur_settings.newaxis: new_axis}}

            yield self.OUTPUT_SIGNAL, replace(msg, data=out_view, dims=out_dims, axes=out_axes)

Let me know if this makes sense and I'll make a PR.