central moment vector

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