diff --git a/src/MultivariateMoments.jl b/src/MultivariateMoments.jl index 4d4bb29..165b0ec 100644 --- a/src/MultivariateMoments.jl +++ b/src/MultivariateMoments.jl @@ -7,7 +7,7 @@ import MutableArithmetics as MA import StarAlgebras as SA import MultivariatePolynomials as MP -const _APL = MP.AbstractPolynomialLike +const _APL = Union{MP.AbstractPolynomialLike,SA.AlgebraElement} import MultivariateBases as MB diff --git a/src/expectation.jl b/src/expectation.jl index d7a1ac6..1e78e3d 100644 --- a/src/expectation.jl +++ b/src/expectation.jl @@ -1,25 +1,43 @@ function _expectation( - μ::MomentVector{S,<:MB.SubBasis{MB.Monomial}}, - p::_APL{T}, + μ::MomentVector{S,<:MB.SubBasis{B}}, + p::SA.AlgebraElement{<:MB.Algebra{BT,B},T}, f, -) where {S,T} +) where {S,T,BT,B} i = firstindex(μ.basis) s = zero(MA.promote_operation(*, S, T)) - for t in MP.terms(p) - while i in eachindex(μ.basis) && MP.monomial(t) != μ.basis.monomials[i] + for (k, v) in SA.nonzero_pairs(SA.coeffs(p)) + mono = SA.basis(p)[k] + while i in eachindex(μ.basis) && mono != μ.basis[i] i += 1 end if !(i in eachindex(μ.basis)) error( - "The polynomial $p has a nonzero term $t with monomial $(MP.monomial(t)) for which the expectation is not known in $μ", + "The polynomial $p has a nonzero term $mono with coefficient $v for which the expectation is not known in $μ", ) end - s += f(μ.values[i], MP.coefficient(t)) + s += f(μ.values[i], v) i += 1 end return s end +function _expectation( + μ::MomentVector{S,<:MB.SubBasis{B}}, + p::SA.AlgebraElement{<:MB.Algebra}, + f, +) where {S,B} + basis = MB.FullBasis{B,MP.monomial_type(typeof(p))}() + return _expectation( + μ, + MB.algebra_element(SA.coeffs(p, basis), basis), + f, + ) +end + +function _expectation(μ::MomentVector, p::MP.AbstractPolynomialLike, f) + return _expectation(μ, MB.algebra_element(MB.sparse_coefficients(MP.polynomial(p)), MB.FullBasis{MB.Monomial,MP.monomial_type(p)}()), f) +end + """ MultivariateMoments.expectation(μ::AbstractMeasureLike, p::AbstractPolynomialLike) MultivariateMoments.expectation(p::AbstractPolynomialLike, μ::AbstractMeasureLike) @@ -37,5 +55,8 @@ expectation(p::_APL, μ::MomentVector) = _expectation(μ, p, (a, b) -> b * a) # See [`expectation`](@ref) """ -LinearAlgebra.dot(μ::AbstractMeasureLike, p::_APL) = expectation(μ, p) -LinearAlgebra.dot(p::_APL, μ::AbstractMeasureLike) = expectation(p, μ) +LinearAlgebra.dot(μ::AbstractMeasureLike, p::MP.AbstractPolynomialLike) = expectation(μ, p) +LinearAlgebra.dot(p::MP.AbstractPolynomialLike, μ::AbstractMeasureLike) = expectation(p, μ) +# Need to split it and not use `_APL` to avoid ambiguity +LinearAlgebra.dot(μ::AbstractMeasureLike, p::SA.AlgebraElement) = expectation(μ, p) +LinearAlgebra.dot(p::SA.AlgebraElement, μ::AbstractMeasureLike) = expectation(p, μ)