Solve PSR using adam

PSR_steady.jl

```
Simulate an unsteady PSR with oscilation
```

using Arrhenius
using LinearAlgebra
using DifferentialEquations
using ForwardDiff
using DiffEqSensitivity
using Sundials
using Plots
using DelimitedFiles
using Profile
using Flux
using Flux: update!
using ForwardDiff: gradient!
using ProgressBars
using Printf
using BSON: @save, @load

# Load cantera data for comparision
cantera_data = readdlm("cantera/data_T.txt");
ct_ts = cantera_data[:, 1];
ct_T = cantera_data[:, 2];
ct_Y = cantera_data[:, 3:end];

# Arrhenius.jl
gas = CreateSolution("./cantera/gri30.yaml");
ns = gas.n_species;

cb = function (p, loss_mean, g_norm; doplot=true)
    global l_loss, l_grad, iter
    push!(l_loss, loss_mean)
    push!(l_grad, g_norm)
    push!(l_pnorm, norm(p))

    if doplot & (iter % n_plot == 0)
        plt_loss = plot(l_loss, yscale=:log10, label="loss");
        plt_grad = plot(l_grad, yscale=:log10, label="grad_norm");
        plt_pnorm = plot(l_pnorm, yscale=:log10, label="p_norm");
        xlabel!(plt_loss, "Epoch");
        xlabel!(plt_grad, "Epoch");
        xlabel!(plt_pnorm, "Epoch");
        ylabel!(plt_loss, "Loss");
        ylabel!(plt_grad, "Grad Norm");
        ylabel!(plt_pnorm, "p Norm");
        plt_all = plot([plt_loss, plt_grad, plt_pnorm]..., legend=:bottomleft);
        png(plt_all, "./figs/loss_grad");

        @save "./checkpoint/mymodel.bson" p opt l_loss l_grad l_pnorm iter
    end
    iter += 1
    return false
end

# Load inlet conditions
# 300, P, 'H2:9.5023,CO:1.7104,CH4:5.7014,O2:17.0090,N2:66.0769'
TYin = readdlm("cantera/TYin.txt")
Yin = zeros(ns)
Tin = 300.0
#
# for (i, s) in enumerate(["H2", "CO", "CH4", "O2", "N2"])
#     Yin[species_index(gas, s)] = TYin[i+1]
# end
#
# Y0 = ct_Y[1, :]
# T0 = ct_T[1]
P = one_atm

T0 = 1500.0
Yin[species_index(gas, "CH4")] = 0.1
Yin[species_index(gas, "O2")] = 0.2
Yin[species_index(gas, "N2")] = 0.7 - 1.e-4
Yin[species_index(gas, "AR")] = 1.e-4
Y0 = Yin

@inbounds function dudt(u, p, t)
    tres = 10.0 # 1.0
    T = u[end]
    Y = @view(u[1:ns])
    mean_MW = 1.0 / dot(Y, 1 ./ gas.MW)
    ρ_mass = P / R / T * mean_MW
    X = Y2X(gas, Y, mean_MW)
    C = Y2C(gas, Y, ρ_mass)
    cp_mole, cp_mass = get_cp(gas, T, X, mean_MW)
    h_mole = get_H(gas, T, Y, X)
    S0 = get_S(gas, T, P, X)
    wdot = wdot_func(gas.reaction, T, C, S0, h_mole)
    Ydot = @. wdot / ρ_mass * gas.MW + (Yin - Y) / tres
    Tdot =
        -dot(h_mole, wdot) / ρ_mass / cp_mass +
        (Tin - T) / tres
    return vcat(Ydot, Tdot)
end

tspan = [0.0, 20.0];
u0 = vcat(Y0, T0);
prob = ODEProblem(dudt, u0, tspan, p);
sol = solve(prob, CVODE_BDF(), reltol = 1e-6, abstol = 1e-9);

u = log.(vcat(sol[1:end-1, end], 1.e0))
p = [2000.0];

function f!(F, x)
    tres = exp(x[end])
    T = 1500.0 * ones(eltype(x), 1)[1]
    Y = exp.(@view(x[1:ns]))
    mean_MW = 1.0 / dot(Y, 1 ./ gas.MW)
    ρ_mass = P / R / T * mean_MW
    X = Y2X(gas, Y, mean_MW)
    C = Y2C(gas, Y, ρ_mass)
    cp_mole, cp_mass = get_cp(gas, T, X, mean_MW)
    h_mole = get_H(gas, T, Y, X)
    S0 = get_S(gas, T, P, X)
    wdot = wdot_func(gas.reaction, T, C, S0, h_mole)
    Ydot = @. wdot / ρ_mass * gas.MW + (Yin - Y) / tres
    Tdot =
        -dot(h_mole, wdot) / ρ_mass / cp_mass +
        (Tin - T) / tres
    F .= vcat(Ydot, Tdot/1000.0)
end
F = similar(u)
f!(F, u)
function j!(J, x)
    F = similar(x)
    ForwardDiff.jacobian!(J, (F, x) -> f!(F, x), F, x)
end

res = nlsolve(f!, j!, u; method = :newton, iterations=1000)

@show sum(exp.(res.zero[1:end-1]))

grad_max = 1.e10;
n_plot = 1000;

l_loss = [];
l_grad = [];
l_pnorm = [];
iter = 1;

grad = similar(u);

opt = ADAMW(0.001, (0.9, 0.999), 1.e-4);
epochs = ProgressBar(1:10000);
for epoch in epochs
    loss = sum(abs.(dudt(u, p)))
    gradient!(grad, u -> sum(abs.(dudt(u, p))), u)
    g_norm = norm(grad)
    if g_norm > grad_max
        @. grad = grad / g_norm * grad_max
    end
    update!(opt, u, grad)
    set_description(epochs, string(@sprintf("loss %.2e pnorm %.2e gnorm %.2e",
                                             loss, norm(u), g_norm)))
    cb(u, loss, g_norm; doplot=true)
end


plt = plot(
    sol.t,
    sol[species_index(gas, "CH4"), :],
    lw = 2,
    label = "Arrhenius.jl",
);
ylabel!(plt, "Mass Fraction of CH4");
xlabel!(plt, "Time [s]");
plt_CO = plot(
    sol.t,
    sol[species_index(gas, "CO"), :],
    lw = 2,
    label = "Arrhenius.jl",
);
ylabel!(plt_CO, "Mass Fraction of CO");
xlabel!(plt_CO, "Time [s]");
pltT = plot(sol.t, sol[end, :], lw = 2, label = "Arrhenius.jl");
ylabel!(pltT, "Temperature [K]");
xlabel!(pltT, "Time [s]");
pltsum = plot(
    plt,
    plt_CO,
    pltT,
    legend = true,
    framestyle = :box,
    layout = (3, 1),
    size = (1200, 1200),
);
png(pltsum, "figs/PSR_SCure.png");