# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 16 Threads 1.6.0
#     language: julia
#     name: julia-16-threads-1.6
# ---

using Random; Random.seed!(1)
using Plots; theme(:gruvbox_dark) # so my eyes don't get seared while working on this at night :)
using LinearAlgebra
using Dierckx
using Revise
using RandomODE

# # Testing the Package
#
# First, we check that the module can correctly create the problem data structure, assemble the linear system to solve the problem, and run the solution of an ensemble of problems correctly:

a = 0; b = 1
nₓ = 100
m(x) = x; σ = 1.; s = 2.
left_bc = DirichletBC(0.); right_bc = DirichletBC(0.)
f(x) = exp(-x)

problem = RandomODEProblem(a, b, nₓ, m, σ, s, left_bc, right_bc, f);

plot(problem.x, problem.m, label = "Mean of Normal Samples", title = "Mean and Right Hand Side Functions")
plot!(problem.x, problem.rhs, label = "f(x)")

potential = sample_field(problem);
problem.V = potential;

plot(problem.x, problem.V, title = "Random Field Sample", label = "V(x)")

assemble_system!(problem);
println("LHS matrix: $(problem.A)")

solve!(problem);
plot(problem.x, problem.u, title = "Computed Solution", label = "u(x)")

# Now we check an ensemble of problems:

# +
Random.seed!(1)
num_problems = 10; problems = []
potentials = []
for prob in 1:num_problems
    problem = RandomODEProblem(a, b, nₓ, m, σ, s, left_bc, right_bc, f)
    potential = sample_field(problem)
    problem.V = potential; push!(potentials, potential)
    push!(problems, problem)
end
potentials_plt = plot(problems[1].x, potentials[1], title = "Potentials", label = "V[1]", legend = :outertopright)
for (i, potential) in enumerate(potentials)
    plot!(problems[i].x, potential, label = "V[$(i)]")
end
display(potentials_plt)

solve_ensemble!(problems)
solutions_plt = plot(problems[1].x, problems[1].u, title = "Solutions", label = "U[1]", legend = :outertopright)
for (i, problem) in enumerate(problems)
    plot!(problems[i].x, problems[i].u, label = "U[$(i)]")
end
display(solutions_plt)
# -

# # Assessment of Convergence

# +
Random.seed!(1)
a = 0; b = 1
nₓ_high = 10000
m(x) = x; σ = 1.; s = 2.
left_bc = DirichletBC(0.); right_bc = DirichletBC(0.)
f(x) = exp(-x)
V(x) = cos(4*π*x)^2

high_res_problem = RandomODEProblem(a, b, nₓ_high, m, σ, s, left_bc, right_bc, f, potential = V);
assemble_system!(high_res_problem)
solve!(high_res_problem)
high_res_spline = Spline1D(high_res_problem.x, high_res_problem.u, k = 1)

nₓ_vals = 100:100:1000
Δx_vals = [(b - a)/(nₓ + 1) for nₓ in nₓ_vals]
errors = []
for nₓ in nₓ_vals
    low_res_problem = RandomODEProblem(a, b, nₓ, m, σ, s, left_bc, right_bc, f, potential = V)
    assemble_system!(low_res_problem)
    solve!(low_res_problem)
    error = norm(low_res_problem.u .- high_res_spline.(low_res_problem.x), Inf)
    push!(errors, error)
end

scatter(Δx_vals, errors, label = "Infinity norm error", 
        title = "Rate of Convergence", xscale = :log10, yscale = :log10, legend = :outertopright)
xlabel!("Δx"); ylabel!("Error")
plot!(Δx_vals, Δx_vals.^2, label = "Quadratic in Δx")
# -

# This clearly displays quadratic convergence in the step size $\Delta x$ for the deterministic problem. This means, that since the convergence rate of Monte Carlo is $\mathcal{O}(M^{-\frac{1}{2}})$, the optimal scaling rate for the mesh is to have $\newline \newline$
# $$
# \Delta x^{2} \propto M^{-\frac{1}{2}},
# \newline
# $$
# or equivalently, $\newline \newline$
# $$
# M \propto \Big(\frac{1}{\Delta x}\Big)^{4}.
# \newline
# $$
#
# Now, we show that there is convergence to a limiting mean as $\sigma \to 0$ and assess convergence with the random potential:

# +
a = 0; b = 1
nₓ = 100
m(x) = x; σ_vals = [1., 0.1, 0.01]
left_bc = DirichletBC(0.); right_bc = DirichletBC(0.)
f(x) = exp(-x)

num_problems = 1e4
first_moments = []; second_moments = []

function run_ensemble(s)
    histogram_plots = []; means_plots = []
    for σ in σ_vals
        problems = []

        # Solve the 10^4 samples with given σ
        for i in 1:num_problems
            problem = RandomODEProblem(a, b, nₓ, m, σ, s, left_bc, right_bc, f)
            potential = sample_field(problem)
            problem.V = potential
            push!(problems, problem)
        end
        solve_ensemble!(problems)

        # Compute first and second moments
        first_moment = zeros(nₓ); second_moment = zeros(nₓ)
        for problem in problems
            first_moment += problem.u
            second_moment += (problem.u).^2
        end
        first_moment /= num_problems; second_moment /= num_problems
        push!(first_moments, first_moment); push!(second_moments, second_moment)

        # Histogram midpoint values
        midpoint_vals = []
        for problem in problems
            problem_spl = Spline1D(problem.x, problem.u, k = 1)
            push!(midpoint_vals, problem_spl(0.5))
        end
        push!(histogram_plots,
            histogram(midpoint_vals, title = "Midpoint values with σ = $(σ)", label = "u(0.5)"))
        
        # Plot mean functions with one standard deviation ribbon
        stddev = @. √(second_moment - first_moment^2)
        means = plot(problems[1].x, first_moment, ribbon=stddev,
                    title = "Mean with One Standard Deviation; σ = $(σ)", label = "Mean")
        push!(means_plots, means)
    end

    # Print maximum pointwise variances
    for (i, σ) in enumerate(σ_vals)
        println("Maximum variance of solutions with σ = $(σ):
                $(maximum(second_moments[i] .- first_moments[i].^2))")
    end
    
    # Plot histograms and mean plots
    hists_plot = plot(histogram_plots[1], histogram_plots[2], histogram_plots[3], layout = (3, 1))
    display(hists_plot)
    means_plot = plot(means_plots[1], means_plots[2], means_plots[3], layout = (3, 1))
    display(means_plot)
end

Random.seed!(1)
run_ensemble(2.)
# -

# It is clear from this output that the maximum pointwise variances are shrinking as $\sigma \to 0$. Therefore, we are converging to a limiting mean. Furthermore, we can look at the histograms of the values at the midpoint, and it is clear that there is convergence to a limiting mean; the tails of the histogram are getting visibly tighter at each smaller value of $\sigma$. Now we repeat the above with $s = 4$:

# +
a = 0; b = 1
nₓ = 100
m(x) = x; σ_vals = [1., 0.1, 0.01]
left_bc = DirichletBC(0.); right_bc = DirichletBC(0.)
f(x) = exp(-x)

num_problems = 1e4
first_moments = []; second_moments = []

Random.seed!(1)
run_ensemble(4.)
# -

# With $s = 4$, the primary difference from the case where $s = 2$ is the smaller variance in solutions. The variance is about an order of magnitude smaller for any particular value of $\sigma$. This is because the higher value of $s$ guarantees a higher degree of smoothness in the potential $V$, and the eigenvalues of the covariance operator are smaller (which in turn reduces the variance of the terms in the Karhunen-Loeve expansion).

StarcoderdataJupyterScriptsDedupFiltered