diff --git a/src/util.jl b/src/util.jl index 00fedb8b0..90b781e34 100644 --- a/src/util.jl +++ b/src/util.jl @@ -4,6 +4,7 @@ import Base: * import LinearAlgebra.BLAS using LinearAlgebra: mul! using FFTW +using Statistics: mean export hilbert, fftintype, @@ -162,11 +163,18 @@ Convert an amplitude ratio to dB (decibel), or ``20 amp2db(a::Real) = 20*log10(a) """ - rms(s) + rms(s; dims) -Return the root mean square of signal `s`. +Return the root mean square (rms) of signal `s`. Optional keyword parameter +`dims` can be used to specify the dimensions along which to compue the rms. """ -rms(s::AbstractArray{T}) where {T<:Number} = sqrt(sum(abs2, s)/length(s)) +function rms(s::AbstractArray{T}; dims=:) where {T<:Number} + if dims === (:) + return sqrt(mean(abs2, s)) + else + return sqrt.(mean(abs2, s; dims=dims)) + end +end """ rmsfft(f) diff --git a/test/util.jl b/test/util.jl index b302a2744..ed6f194e7 100644 --- a/test/util.jl +++ b/test/util.jl @@ -83,6 +83,9 @@ end n = (5,6) for x in ( randn(n), randn(n)+randn(n)im ) @test rms(x) ≈ sqrt(mean(abs.(x).^2)) + @test rms(x; dims=1) ≈ sqrt.(mean(abs.(x).^2; dims=1)) + @test rms(x; dims=2) ≈ sqrt.(mean(abs.(x).^2; dims=2)) + @test rms(x; dims=1:2) ≈ sqrt.(mean(abs.(x).^2; dims=1:2)) @test rmsfft(fft(x)) ≈ rms(x) end # for end