diff --git a/src/koopman.jl b/src/koopman.jl index 7b50d25..b477a21 100644 --- a/src/koopman.jl +++ b/src/koopman.jl @@ -123,35 +123,40 @@ function expectation(g::Function, prob::ODEProblem, u0, p, expalg::MonteCarlo, a end """ - centralmoment(n, g, args...; kwargs) -> [n by 1 Array] + centralmoment(n, g, args...; kwargs) Computes the n central moments of the function g using the Koopman expectation. The function is a wrapper over expectation, arguments can be piped through with args and kwargs. -Return: n-length array of the 1 to n central moments +Return: nout-length array containing the n-length array of the 1 to n central +moments for each dimension of (possibly) multivariate g Note: The first central moment is, by definition, always 0 -TODO: - add support for vector-valued g functions, currently assumes scalar - return values. - - add tests +TODO: - add tests """ -function centralmoment(n::Int, g::Function, args...; kwargs...) - if n < 2 return Float64[] end +function centralmoment(n::Int, g::Function, args...; nout=1, kwargs...) + if n < 2 return nothing end + + # only want to evaluate g once + function g_raisedn(x) + y = g(x) + return ArrayPartition([y.^i for i in 1:n]...) + end # Compute the expectations of g, g^2, ..., g^n - sol = expectation(x -> [g(x)^i for i in 1:n], args...; nout = n, kwargs...) - exp_set = sol.u - mu_g = popfirst!(exp_set) + sol = expectation(g_raisedn, args...; nout = n*nout, kwargs...) + gn_expvals = sol.u # Combine according to binomial expansion - const_term(n) = (-1)^(n-1) * (n-1) * mu_g^n binom_term(n, k, mu, exp_gi) = binomial(n, k) * (-mu)^(n - k) * exp_gi - binom_sum = function (exp_vals) - m = length(exp_vals) + 1 - sum([binom_term(m, k + 1, mu_g, v) for (k,v) in enumerate(exp_vals)]) + const_term(m) + binom_sum = function (i, n, gi_expvals) + mu = gi_expvals.x[1][i] + iterms = [binom_term(n, k, mu, gi_expvals.x[k][i]) for k in 1:n] + return sum(iterms) + (-1)^n * mu^n end + gi_moments = [[0., [binom_sum(i, l, gn_expvals) for l in 2:n]...] for i in 1:nout] - return [zero(exp_set[1]), [binom_sum(exp_set[1:i]) for i in 1:length(exp_set)]...] + return nout == 1 ? gi_moments[1] : gi_moments end