Need help implementing non-allocating normal distribution with StaticArrays

General Usage Performance
2

My first step when unexpected allocations occur is to look at the output of @code_warntype, because often they occur due to type inference issues.

julia> @code_warntype T.log_pdf(comp, x)
MethodInstance for Main.T.log_pdf(::Main.T.MvNormalF, ::SVector{3, Float64})
  from log_pdf(d::Main.T.MvNormalF, x::SVector{3, Float64}) in Main.T at REPL[88]:22
Arguments
  #self#::Core.Const(Main.T.log_pdf)
  d::Main.T.MvNormalF
  x::SVector{3, Float64}
Body::Float64
1 ─ %1  = Main.T.Float64::Core.Const(Float64)
│   %2  = Base.getindex(x, 1)::Float64
│   %3  = Base.getproperty(d, :μ)::SVector{3, Float64}
│   %4  = Base.getindex(%3, 1)::Float64
│   %5  = (%2 - %4)::Float64
│   %6  = Base.getindex(x, 2)::Float64
│   %7  = Base.getproperty(d, :μ)::SVector{3, Float64}
│   %8  = Base.getindex(%7, 2)::Float64
│   %9  = (%6 - %8)::Float64
│   %10 = Base.getindex(x, 3)::Float64
│   %11 = Base.getproperty(d, :μ)::SVector{3, Float64}
│   %12 = Base.getindex(%11, 3)::Float64
│   %13 = (%10 - %12)::Float64
│   %14 = Base.getproperty(d, :Σ_inv)::SMatrix{3, 3, Float64} # <-- marked in red
│   %15 = Main.T.vXvt(%5, %9, %13, %14)::Float64
│   %16 = Base.convert(%1, %15)::Float64
│   %17 = Core.typeassert(%16, %1)::Float64
└──       return %17

It doesn’t show it in the above code, but the line
%14 = Base.getproperty(d, :Σ_inv)::SMatrix{3, 3, Float64} shows an issue.

That is because SMatrix{3, 3, Float64} is not a concrete type.

julia> isconcretetype(SMatrix{3,3,Float64})
false

In fact, SMatrix has a fourth type parameter equal to the total number of elements (see the help entry for more info)

Solution: Change the definition to CovMat = SMatrix{3, 3, Float64, 9} and your code should be type stable and allocation free.

Note that you can construct an MvNormal from Distributions with static arrays too.

Also, it seems there is a bug in your log_pdf

julia> comp = T.MvNormalF(T.MeanVec(0., 0., 0.), T.CovMat(1., 0., 0., 0., 1., 0., 0., 0., 1.))
Main.T.MvNormalF([0.0, 0.0, 0.0], [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0], 2.756815599614018, [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0])

julia> D = MvNormal(comp.μ, comp.Σ)
MvNormal{Float64, PDMats.PDMat{Float64, SMatrix{3, 3, Float64, 9}}, SVector{3, Float64}}(
dim: 3
μ: [0.0, 0.0, 0.0]
Σ: [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]
)

julia> logpdf(D, x) - T.log_pdf(comp, x)
-5.7568155996140185
1 Like