Hi, I’m not sure how to phrase this question, so please bear with me:
I’ve written a function f which takes an input x and returns a vector [p_1(x),...,p_K(x)], which correspond to the first K orthogonal polynomials corresponding to some distribution \mu. I can check that my implementation is right by sampling many X_i \sim \mu and checking that the sample average and covariance matrix are roughly the zero vector and the identity matrix, respectively.
I was wondering if there is a more systematic way of doing this, where I’m willing to restrict myself to distributions in Distributions.jl. My first thought was the following, which I’ll explain informally: I figured you could define a symbolic variable X and pass this to f(X) to get a vector of symbolic polynomials [p_1(x),...,p_K(x)], for which you’d then have to compute the expected value of each polynomial term and simplify; you could get the variance-covariance matrix in a similar way, as this should be the matrix of second moments.
I was wondering if there is any work on this, or how you would suggest a rudimentary implementation should proceed; if nobody is working on this, I’d like to toy around with this a little bit to see if I can get it to work in a simple setting. I guess building on Symbolics.jl and Distributions.jl would be a sensible starting point, and a general flow would look like:
using Symbolics, Distributions
f(t) = x^5 - x + 1
function test(X::Distribution)
mean_f = calculate_expectation(f, X)
end
where calculate_expectation effectively does
function calculate_expectation(f, X)
@variable x
lowered_expression = Symbolics.toexpr(f(x))
# Iterating over the lowered expression, use registered arithmetic rules e.g.
# - expectation of sum is sum of expectations
# - expectation of product is product of expectations under independence
# - expectation of power is moment
# - expectation of exp(xt) is MGF (and similar for other quantities)
# to compute the expectation if possible, returning NaN (or an error)
# if an operation cannot be reduced because no rule is specified
end
I wrote up a very simple proof of concept to compute expectations of polynomials in centered normal random variables. Clearly this needs a bit more careful infrastructure to deal with multivariate expressions and more general random variables, but I’m wondering if this general approach is feasible, or if I should be taking a very different route here?
using Symbolics, Distributions
split_out_sums(expr) = split_out_sums(expr.args[1], expr)
split_out_sums(c::Real) = [c]
split_out_sums(::Any, expr) = [expr]
split_out_sums(::typeof(+), expr) = vcat(split_out_sums(expr.args[2]), split_out_sums(expr.args[3]))
split_out_products(expr) = split_out_products(expr.args[1], expr)
split_out_products(c::Real) = [c]
split_out_products(c::Symbol) = [c]
split_out_products(::Any, expr) = [expr]
split_out_products(::typeof(*), expr) = vcat(split_out_products(expr.args[2]), split_out_products(expr.args[3]))
get_dependent_symbols(expr) = [get_dependent_symbols(expr.args[2]), get_dependent_symbols(expr.args[3])]
get_dependent_symbols(::Real) = Symbol[]
get_dependent_symbols(c::Symbol) = [c]
calculate_expectation_no_products(c::Real, X) = c
calculate_expectation_no_products(c::Symbol, X) = (display(c); error())
function calculate_expectation_no_products(::typeof(getindex), expr, X::Distribution{Univariate, <:ValueSupport}; p::Int = 1)
if expr.args[2] == :x
if expr.args[3] == 1
if p == 1
return mean(X)
else
return get_moments(X, p)
end
else
error("Univariate distribution only has one index")
end
else
error("Unexpected")
end
end
function calculate_expectation_no_products(::typeof(^), expr, X)
if isa(expr.args[2].args[1], typeof(getindex))
calculate_expectation_no_products(expr.args[2].args[1], expr.args[2], X; p = expr.args[3])
else
error("Composite powers have not yet been implemented")
end
end
calculate_expectation_no_products(expr::Expr, X) = calculate_expectation_no_products(expr.args[1], expr, X)
calculate_expectation(c::Real, X) = c
function calculate_expectation(expr::Expr, X)
prod_components = split_out_products(expr)
symbol_vec = [get_dependent_symbols(comp) for comp in prod_components]
# check if all components only refer to different symbols
independent_terms = true
for i in eachindex(symbol_vec), j in i+1:length(symbol_vec)
if length(intersect(symbol_vec[i], symbol_vec[j])) > 0
independent_terms = false
break
end
end
if independent_terms
return prod(calculate_expectation_no_products(comp, X) for comp in prod_components)
else
display("Could not verify that the terms in $prod_components are independent")
error()
end
end
get_n_variates(X::Normal) = 1
get_moments(X::Normal, p::Int) = X.μ == 0 ? (p % 2 == 0 ? X.σ^p * prod((p-1):-2:1) : 0) : error("Not yet implemented")
function calculate_expectation(f::Function, X::Distribution)
d = get_n_variates(X)
@variables x[1:d]
lowered_expression = Symbolics.toexpr(f(x))
# Split out sums
expr_parts = split_out_sums(lowered_expression)
# Operate on each component of the sum individually, then sum
return sum(calculate_expectation(expr, X) for expr in expr_parts)
end
f(x) = 2*x[1]^2 - x[1] + 1
function test()
X = Normal(0, 1)
mean_f = calculate_expectation(f, X)
end
test()
What if you use a formal polynomial instead of a symbolic function ? That would simplify a lot the code as you can extract the coefficients of the polynomial by themselves. And there are tools to convert from symbolic expressions to polynomial IIRC
Because here what would happen if you pass a function that is not a polynomial such as f=sin ?
The reason why I want to work in larger generality is that I’d like to allow non-polynomial functions as well. Of course, whether these will lead to computable situations will depend on the specific case and may only work for some distributions (odd functions and symmetric functions work together; exponential functions will work for functions for which we have the mgf; etc)
If you 1) cover a large part of distributions from Distributions.jl and a large part of potentially known expectations-of-functions, and 2) provide a sensible but slow generic fallback (integrating against the density with quadgk ? A Monte-Carlo approximation? Something smarter?), I would definitively see the proposed functionality as a standalone package.
You could try to interoperate with Distributions.expectation function, undocumented yet: maybe some of the rules you want to store for expectations do belong there instead ?