Support multi-column filtering for DF2TFilter

src/Filters/filt.jl

@@ -117,14 +117,16 @@ filt!(out, f::FilterCoefficients{:z}, x) = filt!(out, convert(SecondOrderSection
 
 """
     DF2TFilter(coef::FilterCoefficients{:z}[, si])
-    DF2TFilter(coef::FilterCoefficients{:z}[, sitype::Type])
+    DF2TFilter(coef::FilterCoefficients{:z}[, sitype::Type][, coldims::Tuple])
 
 Construct a stateful direct form II transposed filter with
 coefficients `coef`.
 
 One can optionally specify as the second argument either
 - `si`, an array representing the initial filter state, or
-- `sitype`, the eltype of a zeroed `si`
+- `sitype`, the eltype of a zeroed `si` and/or `coldims`, the size of extra dimensions
+  of the input. E.g. to column-wise filter an input with dims `(L, N1, N2)`, set
+  `coldims` to `(N1, N2)`.
 
 The initial filter state defaults to zeros if called with one argument.
 
@@ -136,37 +138,40 @@ struct DF2TFilter{T<:FilterCoefficients{:z},S<:Array}
     coef::T
     state::S
 
-    function DF2TFilter{Ti,Si}(coef::PolynomialRatio{:z}, state::Vector) where {Ti,Si}
-        length(state) == max(length(coefa(coef)), length(coefb(coef)))-1 ||
+    function DF2TFilter{Ti,Si}(coef::PolynomialRatio{:z}, state::Array) where {Ti,Si}
+        size(state, 1) == max(length(coefa(coef)), length(coefb(coef)))-1 ||
             throw(ArgumentError("length of state vector must match filter order"))
         new{Ti,Si}(coef, state)
     end
-    function DF2TFilter{Ti,Si}(coef::SecondOrderSections{:z}, state::Matrix) where {Ti,Si}
+    function DF2TFilter{Ti,Si}(coef::SecondOrderSections{:z}, state::Array) where {Ti,Si}
         (size(state, 1) == 2 && size(state, 2) == length(coef.biquads)) ||
             throw(ArgumentError("state must be 2 x nbiquads"))
         new{Ti,Si}(coef, state)
     end
-    function DF2TFilter{Ti,Si}(coef::Biquad{:z}, state::Vector) where {Ti,Si}
-        length(state) == 2 || throw(ArgumentError("length of state must be 2"))
+    function DF2TFilter{Ti,Si}(coef::Biquad{:z}, state::Array) where {Ti,Si}
+        size(state,1) == 2 || throw(ArgumentError("length of state must be 2"))
         new{Ti,Si}(coef, state)
     end
 end
 
-DF2TFilter(coef::Union{PolynomialRatio{:z,T},Biquad{:z,T}}, state::Vector{S}) where {T,S} =
-    DF2TFilter{typeof(coef),Vector{S}}(coef, state)
+DF2TFilter(coef::Union{PolynomialRatio{:z,T},Biquad{:z,T}}, state::S) where {T,S<:Array} =
+    DF2TFilter{typeof(coef),S}(coef, state)
 
-DF2TFilter(coef::SecondOrderSections{:z,T,G}, state::Matrix{S}) where {T,G,S} =
-    DF2TFilter{SecondOrderSections{:z,T,G},Matrix{S}}(coef, state)
+DF2TFilter(coef::SecondOrderSections{:z,T,G}, state::S) where {T,G,S<:Array} =
+    DF2TFilter{SecondOrderSections{:z,T,G},S}(coef, state)
 
 ## PolynomialRatio
-DF2TFilter(coef::PolynomialRatio{:z,T}, ::Type{V}=T) where {T,V} =
-    DF2TFilter(coef, zeros(promote_type(T, V), max(length(coefa(coef)), length(coefb(coef))) - 1))
+DF2TFilter(coef::PolynomialRatio{:z,T}, coldims::Tuple{Vararg{Integer}}) where {T} = DF2TFilter(coef, T, coldims)
+DF2TFilter(coef::PolynomialRatio{:z,T}, ::Type{V}=T, coldims::Tuple{Vararg{Integer}}=()) where {T,V} =
+    DF2TFilter(coef, zeros(promote_type(T, V), max(length(coefa(coef)), length(coefb(coef))) - 1, coldims...))
 
-function filt!(out::AbstractVector, f::DF2TFilter{<:PolynomialRatio,<:Vector}, x::AbstractVector)
-    length(x) != length(out) && throw(ArgumentError("out size must match x"))
+function filt!(out::AbstractArray{<:Any, N}, f::DF2TFilter{<:PolynomialRatio,Array{T,N}} where T, x::AbstractArray{<:Any, N}) where N
+    size(x) != size(out) && throw(ArgumentError("out size must match x"))
 
     si = f.state
-    n = length(si) + 1
+    size(x)[2:end] != size(si)[2:end] && throw(ArgumentError("state size must match x"))
+
+    n = size(si, 1) + 1
     b = coefb(f.coef)
     if n == 1
         mul!(out, x, b[1])
@@ -180,40 +185,54 @@ function filt!(out::AbstractVector, f::DF2TFilter{<:PolynomialRatio,<:Vector}, x
             elseif bs < n
                 append!(b, zero(eltype(b)) for _ in 1:(n-bs))
             end
-            _filt_iir!(out, b, a, x, si, 1)
+            for col in CartesianIndices(axes(x)[2:end])
+                _filt_iir!(out, b, a, x, view(si, :, col), col)
+            end
         elseif n <= SMALL_FILT_CUTOFF
             vtup = ntuple(j -> b[j], Val(length(b)))
-            si .= getfield.(_filt_fir!(out, vtup, x, si, Val(:DF2)), :value)
+            for col in CartesianIndices(axes(x)[2:end])
+                _filt_fir!(out, vtup, x, view(si, :, col), col, Val(true))
+            end
         else
-            _filt_fir!(out, b, x, si, 1)
+            for col in CartesianIndices(axes(x)[2:end])
+                _filt_fir!(out, b, x, view(si, :, col), col)
+            end
         end
     end
     return out
 end
 
 ## SecondOrderSections
-DF2TFilter(coef::SecondOrderSections{:z,T,G}, ::Type{V}=T) where {T,G,V} =
-    DF2TFilter(coef, zeros(promote_type(T, G, V), 2, length(coef.biquads)))
+DF2TFilter(coef::SecondOrderSections{:z,T}, coldims::Tuple{Vararg{Integer}}) where {T} = DF2TFilter(coef, T, coldims)
+DF2TFilter(coef::SecondOrderSections{:z,T,G}, ::Type{V}=T, coldims::Tuple{Vararg{Integer}}=()) where {T,G,V} =
+    DF2TFilter(coef, zeros(promote_type(T, G, V), 2, length(coef.biquads), coldims...))
 
-function filt!(out::AbstractVector, f::DF2TFilter{<:SecondOrderSections,<:Matrix}, x::AbstractVector)
-    length(x) != length(out) && throw(ArgumentError("out size must match x"))
-    _filt!(out, f.state, f.coef, x, 1)
-    out
+function filt!(out::AbstractArray{<:Any, N}, f::DF2TFilter{<:SecondOrderSections,<:Array}, x::AbstractArray{<:Any, N}) where N
+    size(x) != size(out) && throw(ArgumentError("out size must match x"))
+    size(x)[2:end] != size(f.state)[3:end] && throw(ArgumentError("state size must match x"))
+    for col in CartesianIndices(axes(x)[2:end])
+        _filt!(out, view(f.state, :, :, col), f.coef, x, col)
+    end
+    return out
 end
 
 ## Biquad
-DF2TFilter(coef::Biquad{:z,T}, ::Type{V}=T) where {T,V} =
-    DF2TFilter(coef, zeros(promote_type(T, V), 2))
+DF2TFilter(coef::Biquad{:z,T}, coldims::Tuple{Vararg{Integer}}) where {T} = DF2TFilter(coef, T, coldims)
+DF2TFilter(coef::Biquad{:z,T}, ::Type{V}=T, coldims::Tuple{Vararg{Integer}}=()) where {T,V} =
+    DF2TFilter(coef, zeros(promote_type(T, V), 2, coldims...))
 
-function filt!(out::AbstractVector, f::DF2TFilter{<:Biquad,<:Vector}, x::AbstractVector)
-    length(x) != length(out) && throw(ArgumentError("out size must match x"))
+function filt!(out::AbstractArray{<:Any, N}, f::DF2TFilter{<:Biquad,Array{T, N}} where T, x::AbstractArray{<:Any, N}) where N
+    size(x) != size(out) && throw(ArgumentError("out size must match x"))
     si = f.state
-    (si[1], si[2]) = _filt!(out, si[1], si[2], f.coef, x, 1)
-    out
+    size(x)[2:end] != size(si)[2:end] && throw(ArgumentError("state size must match x"))
+    for col in CartesianIndices(axes(x)[2:end])
+        (si[1, col], si[2, col]) = _filt!(out, si[1, col], si[2, col], f.coef, x, col)
+    end
+    return out
 end
 
 """
-    filt(f::DF2TFilter{<:FilterCoefficients{:z},<:Array{T}}, x::AbstractVector{V}) where {T,V}
+    filt(f::DF2TFilter{<:FilterCoefficients{:z},<:Array{T}}, x::AbstractArray{V}) where {T,V}
 
 Apply the [stateful filter](@ref stateful-filter-objects) `f` on `x`.
 
@@ -223,13 +242,14 @@ Apply the [stateful filter](@ref stateful-filter-objects) `f` on `x`.
     For more control over the output type, provide a preallocated
     output array `out` to `filt!(out, f, x)`.
 """
-filt(f::DF2TFilter{<:FilterCoefficients{:z},<:Array{T}}, x::AbstractVector{V}) where {T,V} =
+filt(f::DF2TFilter{<:FilterCoefficients{:z},<:Array{T}}, x::AbstractArray{V}) where {T,V} =
     filt!(similar(x, promote_type(T, V)), f, x)
 
 # Fall back to SecondOrderSections
-DF2TFilter(coef::FilterCoefficients{:z}) = DF2TFilter(convert(SecondOrderSections, coef))
-DF2TFilter(coef::FilterCoefficients{:z}, arg::Union{Matrix,Type}) =
-    DF2TFilter(convert(SecondOrderSections, coef), arg)
+DF2TFilter(coef::FilterCoefficients{:z}, coldims::Tuple{Vararg{Integer}}=()) =
+    DF2TFilter(convert(SecondOrderSections, coef), coldims)
+DF2TFilter(coef::FilterCoefficients{:z}, ::Type{V}, coldims::Tuple{Vararg{Integer}}=()) where {V} =
+    DF2TFilter(convert(SecondOrderSections, coef), V, coldims)
 
 #
 # filtfilt

src/dspbase.jl

@@ -118,12 +118,11 @@ end
 const SMALL_FILT_VECT_CUTOFF = 19
 
 # Transposed direct form II
-@generated function _filt_fir!(out, b::NTuple{N,T}, x, siarr, colv) where {N,T}
+@generated function _filt_fir!(out, b::NTuple{N,T}, x, siarr, col, ::Val{StoreSI}=Val(false)) where {N,T,StoreSI}
     silen = N - 1
     si_end = Symbol(:si_, silen)
 
     quote
-        col = colv isa Val{:DF2} ? CartesianIndex() : colv
         N <= SMALL_FILT_VECT_CUTOFF && checkbounds(siarr, $silen)
         Base.@nextract $silen si siarr
         for i in axes(x, 1)
@@ -137,9 +136,10 @@ const SMALL_FILT_VECT_CUTOFF = 19
                 out[i, col] = val
             end
         end
-        if colv isa Val{:DF2}
-            return Base.@ntuple $silen j -> VecElement(si_j)
+        if StoreSI
+            Base.@nexprs $silen j -> siarr[j] = si_j
         end
+        return nothing
     end
 end
 

test/filt.jl

@@ -63,6 +63,9 @@ using DSP, Test, Random, FilterTestHelpers
     df = digitalfilter(Lowpass(0.25), Butterworth(4))
     f = @test_nowarn DF2TFilter(df, ComplexF64)
     @test_nowarn filt(f, s)
+
+    # DF2TFilter{ZPG/SOS} stackoverflow error
+    @test_throws MethodError DF2TFilter(ZeroPoleGain([1], [2], 3), :D, 'x', Ref(1))
 end
 
 @testset "multi-column filt $D-D" for D in 1:4
@@ -89,6 +92,26 @@ end
      @test all(col -> col ≈ y_ref, eachslice(filt(PolynomialRatio(b, a), x); dims=slicedims))
 end
 
+@testset "multi-column DF2TFilter $D-D" for D in 1:4
+    b = [0.1, 0.1]
+    a = [1.0, -0.8]
+    sz = (10, ntuple(n -> n+1, Val(D))...)
+    y_ref = filt(b, a, ones(2*sz[1]))
+    x = ones(sz)
+
+    H = DF2TFilter(PolynomialRatio(b, a),sz[2:end])
+    @test all(col -> col ≈ y_ref[1:sz[1]], eachslice(filt(H, x); dims=ntuple(n -> n+1, Val(D))))
+    @test all(col -> col ≈ y_ref[sz[1]+1:end], eachslice(filt(H, x); dims=ntuple(n -> n+1, Val(D))))
+
+    H = DF2TFilter(SecondOrderSections(PolynomialRatio(b, a)), sz[2:end])
+    @test all(col -> col ≈ y_ref[1:sz[1]], eachslice(filt(H, x); dims=ntuple(n -> n+1, Val(D))))
+    @test all(col -> col ≈ y_ref[sz[1]+1:end], eachslice(filt(H, x); dims=ntuple(n -> n+1, Val(D))))
+
+    H = DF2TFilter(Biquad(PolynomialRatio(b, a)), sz[2:end])
+    @test all(col -> col ≈ y_ref[1:sz[1]], eachslice(filt(H, x); dims=ntuple(n -> n+1, Val(D))))
+    @test all(col -> col ≈ y_ref[sz[1]+1:end], eachslice(filt(H, x); dims=ntuple(n -> n+1, Val(D))))
+end
+
 #
 # filtfilt
 #