Support StarAlgebras in expectation

JuliaAlgebra · Jun 10, 2024 · 3b141ea · 3b141ea
1 parent 71944e4
commit 3b141ea
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 10 deletions.
diff --git 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
@@ -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, μ)