Stochastic Galerkin for ODESystem and PDESystem

.github/workflows/CI.yml

@@ -12,7 +12,7 @@ jobs:
           - "1"
           - "1.6"
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}

Project.toml

@@ -4,22 +4,24 @@ authors = ["Bowen S. Zhu <[email protected]> and contributors"]
 version = "0.1.1"
 
 [deps]
+Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+PolyChaos = "8d666b04-775d-5f6e-b778-5ac7c70f65a3"
 RandomizedLinAlg = "0448d7d9-159c-5637-8537-fd72090fea46"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 TSVD = "9449cd9e-2762-5aa3-a617-5413e99d722e"
 
 [compat]
+Combinatorics = "1.0"
 DocStringExtensions = "0.8, 0.9"
 ModelingToolkit = "8.21"
+PolyChaos = "0.2"
 RandomizedLinAlg = "0.1"
 Setfield = "0.8, 1"
-SymbolicUtils = "0.19"
 Symbolics = "4.10"
 TSVD = "0.4"
 julia = "1.6"

docs/Project.toml

@@ -6,6 +6,8 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+PolyChaos = "8d666b04-775d-5f6e-b778-5ac7c70f65a3"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
 Documenter = "0.27"

docs/make.jl

@@ -1,4 +1,4 @@
-using Documenter, ModelOrderReduction
+using Documenter, ModelOrderReduction, Symbolics, PolyChaos
 
 include("pages.jl")
 

docs/pages.jl

@@ -1,5 +1,12 @@
 pages = [
     "Home" => "index.md",
-    "Functions" => "functions.md",
-    "Tutorials" => Any["tutorials/deim_FitzHugh-Nagumo.md"],
+    "Polynomial Chaos Expansion (PCE)" => [
+        "Introduction" => "pce/pce.md",
+        "pce/stochastic_galerkin.md",
+    ],
+    "Discrete Empirical Interpolation Method (DEIM)" => [
+        "Introduction" => "deim/deim.md",
+        "deim/deim_FitzHugh-Nagumo.md",
+    ],
+    "internals.md",
 ]

docs/src/deim/deim.md

@@ -0,0 +1,5 @@
+# Discrete Empirical Interpolation Method (DEIM)
+
+```@docs
+deim
+```

docs/src/tutorials/deim_FitzHugh-Nagumo.md → docs/src/deim/deim_FitzHugh-Nagumo.md

@@ -1,4 +1,4 @@
-# Discrete Empirical Interpolation Method (DEIM)
+# Example: FitzHugh-Nagumo System
 
 This section illustrates how ModelOrderReduction.jl can be used to build a reduced order 
 model via Petrov-Galerkin projection using the Proper Orthogonal Decomposition (POD) and 

docs/src/internals.md

@@ -0,0 +1,6 @@
+# Internal Documentation
+
+```@autodocs
+Modules = [ModelOrderReduction]
+Public = false
+```

docs/src/pce/pce.md

@@ -0,0 +1,6 @@
+# Polynomial Chaos Expansion (PCE)
+
+```@docs
+PCE
+PCE(::AbstractVector{Num}, ::AbstractVector{Num}, ::AbstractVector{Num}, ::AbstractVector{<:Union{AbstractOrthoPoly, AbstractCanonicalOrthoPoly}})
+```

docs/src/pce/stochastic_galerkin.md

@@ -0,0 +1,5 @@
+# Stochastic Galerkin Method (SGM)
+
+```@docs
+stochastic_galerkin
+```

examples/PCE/tutorial.jl

@@ -0,0 +1,123 @@
+using ModelOrderReduction, ModelingToolkit, DifferentialEquations,
+      PolyChaos, Plots, LinearAlgebra
+
+# Reaction system
+# k[1](1+θ[1]/2)/k[2]: A + B ⇌ C 
+# k[3](1+θ[2]/2):      C → D
+# k[4]:                B → D
+# 
+# state: c[1:4] = [A, B, C, D]
+# certain parameters: k[1:4]
+# uncertain parameters: θ[1:2] ∼ U[-1,1] 
+
+@parameters k[1:4], θ[1:2]
+@variables t, c(t)[1:4]
+D = Differential(t)
+reactor = [D(c[1]) ~ -k[1] * (1 + 0.5 * θ[1]) * c[1] * c[2] + k[2] * c[3];
+           D(c[2]) ~ -k[1] * (1 + 0.5 * θ[1]) * c[1] * c[2] - k[4] * c[2] + k[2] * c[3];
+           D(c[3]) ~ k[1] * c[1] * c[2] - k[3] * c[3];
+           D(c[4]) ~ k[3] * (1 + 0.5 * θ[2]) * c[3] + k[4] * c[2]];
+
+@named reactor_model = ODESystem(reactor, t, c, vcat(k, θ))
+
+# 1. choose/generate PCE of the system state
+d_pce = 3
+pce = PCE(c, [θ[i] => Uniform_11OrthoPoly(d_pce) for i in eachindex(θ)])
+
+# 2. generate moment equations
+moment_eqs, pce_eval = moment_equations(reactor_model, pce)
+
+# 3. solve the moment equations
+# 3a. find initial moments via Galerkin projection (note that c0 could depend on an uncertain parameter)
+c0 = [1.0, 2.0, 0.0, 0.0]
+z0 = reduce(vcat, pce_galerkin(c0, pce))
+
+# 3b. solve the initial value problem for the moment equations
+moment_prob = ODEProblem(moment_eqs, z0, (0.0, 10.0),
+                         [k[1] => 1.0, k[2] => 0.2, k[3] => 12, k[4] => 0.1])
+moment_sol = solve(moment_prob, Tsit5())
+
+# 3c. define a function approximating the parametric solution to the IVP.
+pce_sol(t, θ) = pce_eval(moment_sol(t), θ)
+
+# 4. now let's compare to true solution
+# take 100 uniform samples
+n = 5
+θsamples = Iterators.product(range(-1, 1, length = n), range(-1, 1, length = n))
+ps = [k[1] => 1.0, k[2] => 0.2, k[3] => 12, k[4] => 0.1, θ[1] => 0.0, θ[2] => 0.0]
+pvals = [p[2] for p in ps]
+reactor_problem = ODEProblem(reactor_model, c0, (0.0, 10.0), ps)
+t_range = range(0.0, 10.0, length = 100)
+sols = []
+for θval in θsamples
+    pvals[end - 1] = θval[1]
+    pvals[end] = θval[2]
+    _prob = remake(reactor_problem, p = pvals)
+    push!(sols, Array(solve(_prob, Tsit5(), saveat = t_range)))
+end
+
+species = ["A", "B", "C", "D"]
+color = [:red, :green, :blue, :orange]
+plots = [plot(title = "Species $(species[i])", xlabel = "time", ylabel = "concentration",
+              legend = false) for i in 1:4]
+for sol in sols
+    for i in 1:4
+        plot!(plots[i], t_range, sol[i, :], color = color[i])
+    end
+end
+for θval in θsamples
+    pce_predictions = [pce_sol(t, [θval...]) for t in t_range]
+    for i in 1:4
+        plot!(plots[i], t_range, [c[i] for c in pce_predictions], color = "black",
+              linestyle = :dash)
+    end
+end
+
+fig = plot(plots...)
+savefig(fig, string(@__DIR__, "/traces.png"))
+
+n = 100
+θsamples = Iterators.product(range(-1, 1, length = n), range(-1, 1, length = n))
+ps = [k[1] => 1.0, k[2] => 0.2, k[3] => 12, k[4] => 0.1, θ[1] => 0.0, θ[2] => 0.0]
+pvals = [p[2] for p in ps]
+reactor_problem = ODEProblem(reactor_model, c0, (0.0, 10.0), ps)
+t_range = range(0.0, 10.0, length = 100)
+sols = []
+for θval in θsamples
+    pvals[end - 1] = θval[1]
+    pvals[end] = θval[2]
+    _prob = remake(reactor_problem, p = pvals)
+    push!(sols, Array(solve(_prob, Tsit5(), saveat = t_range)))
+end
+
+mean_solution = mean(sols)
+var_solution = mean([sol .^ 2 for sol in sols]) - mean_solution .^ 2
+
+L = length(pce.sym_basis)
+mean_PCE = [[moment_sol(t)[((i - 1) * L + 1):(i * L)][1] for i in 1:4] for t in t_range]
+var_weightings = computeSP2(pce.tensor_basis)
+var_PCE = [[dot(moment_sol(t)[((i - 1) * L + 1):(i * L)] .^ 2, var_weightings) for i in 1:4]
+           for t in t_range] .- [m .^ 2 for m in mean_PCE]
+
+species = ["A", "B", "C", "D"]
+color = [:red, :green, :blue, :orange]
+plots = [plot(title = "Species $(species[i])", xlabel = "time", ylabel = "mean",
+              legend = false) for i in 1:4]
+
+for i in 1:4
+    plot!(plots[i], t_range, [mean_solution[i, k] for k in axes(mean_solution, 2)],
+          color = color[i])
+    plot!(plots[i], t_range, [s[i] for s in mean_PCE], color = "black", linestyle = :dash)
+end
+fig = plot(plots...)
+savefig(fig, string(@__DIR__, "/mean.png"))
+
+plots = [plot(title = "Species $(species[i])", xlabel = "time", ylabel = "variance",
+              legend = false) for i in 1:4]
+for i in 1:4
+    plot!(plots[i], t_range, [var_solution[i, k] for k in axes(var_solution, 2)],
+          color = color[i])
+    plot!(plots[i], t_range, [s[i] for s in var_PCE], color = "black", linestyle = :dash)
+end
+fig = plot(plots...)
+savefig(fig, string(@__DIR__, "/var.png"))

examples/PCE/tutorial.md

@@ -0,0 +1,109 @@
+# Polynomial Chaos Expansion for Random ODEs
+## Sketch of the Theory
+This section reviews the basic concept behind the use of Polynomial Chaos Expansion (PCE) for approximate solution to random ODEs as well as how `ModelOrderReduction` facilitates its use within the SciML ecosystem.
+
+In essence, PCE is nothing more than a function expansion technique for square integrable functions where the basis functions are chosen to be orthogonal polynomials; for a detailed treatment of the mathematical underpinnings the interested reader is referred to Gautschi (2004) [1], for example. For the sake of this tutorial we will remain much more informal and discuss only aspects that are critical for basic use of PCE in computational practice. To that end, we will consider the problem of finding a parametric solution to the following *random* ODE:
+```math
+\begin{aligned}
+    \frac{dx}{dt}(p) = f(x,p,t), \quad x(0) = x_0(p)
+\end{aligned}
+```
+where the parameters $p$ are *apriori* only known to follow a given distribution $\mu$. Specifically, we will seek to find an approximation to the solution $x(t,p)$ in the form of a PCE that matches the statistics of the random process $x(t,p)$ with $p \sim \mu$. To that end, we approximate 
+```math
+\begin{aligned}
+    x(t,p) \approx \hat{x}(t,p) = \sum_{i=0}^L z_i(t) \zeta_i(p)
+\end{aligned}
+```
+where $\zeta_0(p), \dots, \zeta_L(p)$ are mutually orthogonal polynomial with respect to the distribution $\mu$, i.e.,
+```math
+\begin{aligned}
+    \left\langle \zeta_i, \zeta_j \right\rangle := \int \zeta_i(p) \zeta_j(p) d\mu(p) = \delta_{ij}
+\end{aligned}
+```
+where $\delta_{ij}$ is the Kronecker delta. These polynomials can be efficiently constructed for many distributions. In particular, they are known analytically for many standard distributions including but not limited to Gaussians, uniform or beta distributions; moreover, they can be approximated numerically for distributions with known smooth densities and even empirically known distributions. For more details, please review the documentation of [PolyChaos.jl](github.com/SciML/PolyChaos.jl) or Gautschi (2004) [1] for instance.
+
+Inserting the PCE-Ansatz $\hat{x}$ into the random ODE of interest yields 
+```math
+\begin{aligned}
+    \sum_{i=0}^L \frac{dz_i}{dt}(t) \zeta_i(p) = f\left(\sum_{i=0}^L z_i(t) \zeta_i(p), p, t \right),  \sum_{i=0}^L z_i(0) \zeta_i(p)
+\end{aligned}
+```
+To derive evolution equations for the coefficients (or *moments*) $z_i$ we simply demand that the residual of the above equations to be orthogonal to the space spanned by $\zeta_0, \dots, \zeta_L$. Orthogonality here shall be understood with respect to the inner product as described above. This is equivalent to demanding minimal expected mismatch when evaluating the above equation for $p \sim \mu$. Mathematically, this condition translates into a set of deterministic ODEs for the moments $z_j$:
+```math
+\begin{aligned}
+    \left\langle\sum_{i=0}^L \frac{dz_i}{dt}(t) \zeta_i(p), \zeta_j(p) \right\rangle = \left \langle f\left(\sum_{i=0}^L z_i(t) \zeta_i(p), p, t \right),  \sum_{i=0}^L z_i(0) \zeta_i(p), \zeta_j(p) \right\rangle
+\end{aligned}
+```
+and since $\left\langle \zeta_i, \zeta_j\right\rangle = \delta_{ij}$ it follows that
+```math
+\begin{aligned}
+    \frac{dz_0}{dt} &= \left \langle f\left(\sum_{i=0}^L z_i(t) \zeta_i(p), p, t \right),  \sum_{i=0}^L z_i(0) \zeta_i(p), \zeta_0(p) \right\rangle\\
+    & \vdots\\
+    \frac{dz_L}{dt} &= \left \langle f\left(\sum_{i=0}^L z_i(t) \zeta_i(p), p, t \right),  \sum_{i=0}^L z_i(0) \zeta_i(p), \zeta_L(p) \right\rangle
+\end{aligned}
+```
+with initial condition
+```math
+\begin{aligned}
+    z_i(0) = \left\langle x_0, \zeta_i \right\rangle. 
+\end{aligned}
+```
+These equations are referred to as the moment equations and can in principle be solved with any suitable IVP solution technique. 
+
+## Example
+In this section, we will showcase how `ModelOrderReduction.jl` may be used to apply PCE for the approximate solution of random ODEs as described in the previous section. To that end, we will consider a simple nonlinear reactor in which the following reaction occurs
+```math
+\begin{aligned}
+    &A + B \xleftrightharpoons[k_1(1-\frac{\theta_1}{2})]{k_2} C\\
+    &C \xrightarrow{k_3(1-\frac{\theta_2}{2})} D\\
+    &B \xrightarrow{k_4} D
+\end{aligned}
+```
+where the rate parameters $\theta_1 \sim U[-1,1]$, $\theta_2 \sim U[-1,1]$ are uncertain. Using [ModelingToolkit.jl]{github.com/SciML/ModelingToolkit.jl} the following code creates the associated model.
+```
+using ModelOrderReduction, ModelingToolkit, DifferentialEquations, PolyChaos, Plots
+
+@parameters k[1:4], θ[1:2]
+@variables t, c(t)[1:4]
+D = Differential(t)
+reactor = [D(c[1]) ~ -k[1]*(1+0.5*θ[1])*c[1]*c[2] + k[2]*c[3];
+           D(c[2]) ~ -k[1]*(1+0.5*θ[1])*c[1]*c[2] - k[4]*c[2] + k[2]*c[3];
+           D(c[3]) ~ k[1]*c[1]*c[2] - k[3]*c[3];
+           D(c[4]) ~ k[3]*(1+0.5*θ[2])*c[3] + k[4]*c[2]];
+
+@named reactor_model = ODESystem(reactor, t, c, vcat(k, θ))
+```
+
+Next, we are going to define the PCE Ansatz with the help of which we are going to approximate the parametric solution. To that end, we heavily rely on [PolyChaos.jl]{github.com/SciML/PolyChaos.jl}. In this example, we choose a PCE of degree 3. 
+```
+d_pce = 3
+pce = PCE(c, [θ[i] => Uniform_11OrthoPoly(d_pce) for i in eachindex(θ)])
+```
+
+Next, we are constructing the corresponding moment equations via a simple function call of `moment_equations`. This function returns an `ModelingToolkit.ODESystem` describing the deterministic evolution equation of the PCE moments and a function that maps the state of this model via the PCE Ansatz back to the original model state (in this example, the concentration of the different species).
+```
+moment_eqs, pce_eval = moment_equations(reactor_model, pce)
+```
+
+Next, we are solving the evolution equation for the PCE moments. The initial condition for the moments is also obtained via Galerkin projection onto the space spanned by the basis functions in the PCE Ansatz. 
+```
+c0 = [1.0, 2.0, 0.0, 0.0]
+z0 = reduce(vcat, pce_galerkin(c0, pce))
+moment_prob = ODEProblem(moment_eqs, z0, (0.0, 10.0), [k[1] => 1.0, k[2] => 0.2, k[3] => 12, k[4] => 0.1])
+moment_sol = solve(moment_prob, Tsit5())
+```
+
+Lastly, we wish to demonstrate that this approximation can be remarkably accurate in practice. The figure below shows the PCE approximation (dashed black lines) compared to the full model for several realization of the uncertain paramters. 
+
+![traces](traces.png)
+
+While PCE can even provide a good parametric surrogate model of the true solution of the random ODE, it is really designed to capture the statistics of the solution of the random ODE well. Moreover, any statistics that can be computed in terms of the moments of the states of the random ODE are remarkably cheap to compute from the PCE coefficients/moments alone. The figures below shows an example for the mean and variances of the different species.
+
+
+![mean](mean.png)
+![var](var.png)
+
+
+
+## References
+Gautschi, Walter. Orthogonal polynomials: computation and approximation. OUP Oxford, 2004.

src/DataReduction/POD.jl

@@ -1,6 +1,3 @@
-using TSVD: tsvd
-using RandomizedLinAlg: rsvd
-
 function matricize(VoV::Vector{Vector{T}}) where {T}
     reduce(hcat, VoV)
 end

src/ModelOrderReduction.jl

@@ -1,23 +1,31 @@
 module ModelOrderReduction
 
+using Combinatorics: powerset, combinations
 using DocStringExtensions
-
-using Symbolics
-using ModelingToolkit
 using LinearAlgebra
-
+using ModelingToolkit
+using PolyChaos
+using RandomizedLinAlg: rsvd
 using Setfield
+using SparseArrays
+using Symbolics
+using TSVD: tsvd
 
 include("utils.jl")
 
 include("Types.jl")
 include("ErrorHandle.jl")
-
 include("DataReduction/POD.jl")
-export SVD, TSVD, RSVD
-export POD, reduce!
+export POD, reduce!, matricize
 
+using Setfield
 include("deim.jl")
 export deim
 
+include("pce/pce.jl")
+export PCE
+include("pce/pce_metadata.jl")
+include("pce/stochastic_galerkin.jl")
+export stochastic_galerkin
+
 end