SciPy lfilter for arbitrary filter, with initial conditions applied along any axis of N-D array-CodePudding

This python numpy/Scipy question SciPy lfilter with initial conditions applied along any axis of N-D array has an answer that works well when the operation axis!=0, or when the length of the initial condition vector zi remains 1, ie for a first-order filter. It fails as soon as a higher order filter is defined to operate on an input axis assigned to the first axis of signal, where signal is 2D.

For example:

import numpy as np
from scipy.signal import (lfilter, lfilter_zi, butter)

def apply_filter(B, A, signal, axis=-1):
   # apply filter, setting proper initial state (doesn't assume rest)
   filtered, zf = lfilter(B, A, signal, 
             zi= lfilter_zi(B, A) * np.take(signal, [0], axis=axis), axis=axis)
   return filtered

B, A = butter(1, 0.5)
x = np.random.randn(12, 50)
apply_filter(B, A, x, axis=1)    # works without error
apply_filter(B, A, x, axis=0)    # works without error


B, A = butter(2, 0.5)
x = np.random.randn(12, 50)
apply_filter(B, A, x, axis=1)    # works without error
apply_filter(B, A, x, axis=0)    # raises error

raises

ValueError: operands could not be broadcast together with shapes (2,) (1,50)

How could the solution be made more general to apply to an arbitrary filter length for axis=0?

CodePudding user response：

Your error came from the shape of your arrays when * these two array. Simple example:

>>> np.random.rand(2) * np.random.rand(1,10) 
----> 1 np.random.rand(2) * np.random.rand(1,10)

ValueError: operands could not be broadcast together with shapes (2,) (1,10) 


>>> np.random.rand(2)[:,None] * np.random.rand(1,10) 
array([[0.05905608, 0.30028617, 0.12495555, 0.28012041, 0.15031258,
        0.05166653, 0.2035891 , 0.01499304, 0.31749996, 0.3146938 ],
       [0.06860488, 0.34883958, 0.14515967, 0.3254132 , 0.17461669,
        0.06002052, 0.23650752, 0.01741727, 0.36883668, 0.36557679]])

For handling different axis in your function you can write the function like the below:

def apply_filter(B, A, signal, axis=-1):
    tmp = lfilter_zi(B, A) if axis==1 else lfilter_zi(B, A)[:,None]
    filtered, zf = lfilter(B, A, signal, zi = tmp * np.take(signal, [0], axis=axis), axis=axis)
    return filtered