Skip to content

Commit

Permalink
test: remove Flux based tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@sathvikbhagavan
sathvikbhagavan committed Jan 23, 2024
1 parent a8492f7 commit 0bd8e1f
Show file tree
Hide file tree
Showing 14 changed files with 205 additions and 1,084 deletions.
42 changes: 5 additions & 37 deletions test/BPINN_PDE_tests.jl
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
using Test, MCMCChains, Lux, ModelingToolkit
import ModelingToolkit: Interval, infimum, supremum
using ForwardDiff, Distributions, OrdinaryDiffEq
using Flux, AdvancedHMC, Statistics, Random, Functors
using AdvancedHMC, Statistics, Random, Functors
using NeuralPDE, MonteCarloMeasurements
using ComponentArrays

Expand All @@ -17,8 +17,6 @@ eqs = Dt(u(t)) - cos(2 * π * t) ~ 0
bcs = [u(0) ~ 0.0]
domains = [t ∈ Interval(0.0, 2.0)]

chainf = Flux.Chain(Flux.Dense(1, 6, tanh), Flux.Dense(6, 1)) |> Flux.f64
init1, re1 = Flux.destructure(chainf)
chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 1))
initl, st = Lux.setup(Random.default_rng(), chainl)

Expand All @@ -35,28 +33,13 @@ sol1 = ahmc_bayesian_pinn_pde(pde_system,
priorsNNw = (0.0, 1.0),
saveats = [1 / 50.0])

discretization = NeuralPDE.BayesianPINN([chainf], GridTraining([0.01]))
sol2 = ahmc_bayesian_pinn_pde(pde_system,
discretization;
draw_samples = 1500,
bcstd = [0.01],
phystd = [0.005],
priorsNNw = (0.0, 1.0),
saveats = [1 / 50.0])

analytic_sol_func(u0, t) = u0 + sin(2 * π * t) / (2 * π)
ts = vec(sol1.timepoints[1])
u_real = [analytic_sol_func(0.0, t) for t in ts]
u_predict = pmean(sol1.ensemblesol[1])
@test u_predict≈u_real atol=0.5
@test mean(u_predict .- u_real) < 0.1

ts = vec(sol2.timepoints[1])
u_real = [analytic_sol_func(0.0, t) for t in ts]
u_predict = pmean(sol2.ensemblesol[1])
@test u_predict≈u_real atol=0.5
@test mean(u_predict .- u_real) < 0.1

## Example 1, 1D ode
@parameters θ
@variables u(..)
Expand All @@ -73,10 +56,9 @@ bcs = [u(0.0) ~ 1.0]
domains = [θ ∈ Interval(0.0, 1.0)]

# Neural network
chain = Lux.Chain(Lux.Dense(1, 12, Flux.σ), Lux.Dense(12, 1))
chain = Lux.Chain(Lux.Dense(1, 12, Lux.σ), Lux.Dense(12, 1))

discretization = NeuralPDE.BayesianPINN([chain],
GridTraining([0.01]))
discretization = BayesianPINN([chain], GridTraining([0.01]))

@named pde_system = PDESystem(eq, bcs, domains, [θ], [u])

Expand Down Expand Up @@ -123,7 +105,7 @@ chain = [
Lux.Dense(10, 1)), Lux.Chain(Lux.Dense(1, 4, Lux.tanh), Lux.Dense(4, 1)),
Lux.Chain(Lux.Dense(1, 4, Lux.tanh), Lux.Dense(4, 1))]

discretization = NeuralPDE.BayesianPINN(chain, GridTraining(0.01))
discretization = BayesianPINN(chain, GridTraining(0.01))

@named pde_system = PDESystem(eq, bcs, domains, [x],
[u(x), Dxu(x), Dxxu(x), O1(x), O2(x)])
Expand All @@ -143,11 +125,6 @@ xs = vec(sol1.timepoints[1])
u_real = [analytic_sol_func(x) for x in xs]
@test u_predict≈u_real atol=0.5

# diff_u = abs.(u_real .- u_predict)
# plot(xs, u_real)
# plot!(xs, u_predict)
# plot!(xs, diff_u)

# 2D Poissons equation
@parameters x y
@variables u(..)
Expand All @@ -171,7 +148,7 @@ chain = Lux.Chain(Lux.Dense(dim, 9, Lux.σ), Lux.Dense(9, 9, Lux.σ), Lux.Dense(

# Discretization
dx = 0.05
discretization = NeuralPDE.BayesianPINN([chain], GridTraining(dx))
discretization = BayesianPINN([chain], GridTraining(dx))

@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])

Expand All @@ -190,12 +167,3 @@ u_predict = pmean(sol1.ensemblesol[1])
u_real = [analytic_sol_func(xs[:, i][1], xs[:, i][2]) for i in 1:length(xs[1, :])]
diff_u = abs.(u_predict .- u_real)
@test u_predict≈u_real atol=1.5

# using Plots, StatsPlots
# plotly()
# plot(sol1.timepoints[1][1, :],
# sol1.timepoints[1][2, :],
# pmean(sol1.ensemblesol[1]),
# linetype = :contourf)
# plot(sol1.timepoints[1][1, :], sol1.timepoints[1][2, :], u_real, linetype = :contourf)
# plot(sol1.timepoints[1][1, :], sol1.timepoints[1][2, :], diff_u, linetype = :contourf)
64 changes: 12 additions & 52 deletions test/BPINN_PDEinvsol_tests.jl
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
using Test, MCMCChains, Lux, ModelingToolkit
import ModelingToolkit: Interval, infimum, supremum
using ForwardDiff, Distributions, OrdinaryDiffEq
using Flux, AdvancedHMC, Statistics, Random, Functors
using AdvancedHMC, Statistics, Random, Functors
using NeuralPDE, MonteCarloMeasurements
using ComponentArrays

Expand All @@ -17,8 +17,6 @@ eqs = Dt(u(t)) - cos(p * t) ~ 0
bcs = [u(0) ~ 0.0]
domains = [t ∈ Interval(0.0, 2.0)]

chainf = Flux.Chain(Flux.Dense(1, 6, tanh), Flux.Dense(6, 1)) |> Flux.f64
init1, re1 = Flux.destructure(chainf)
chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 1))
initl, st = Lux.setup(Random.default_rng(), chainl)

Expand All @@ -36,13 +34,9 @@ u = [analytic_sol_func1(0.0, timepoint) for timepoint in timepoints]
u = u .+ (u .* 0.2) .* randn(size(u))
dataset = [hcat(u, timepoints)]

# plot(dataset[1][:, 2], dataset[1][:, 1])
# plot!(timepoints, u)

# checking all training strategies
discretization = NeuralPDE.BayesianPINN([chainl],
StochasticTraining(200),
param_estim = true, dataset = [dataset, nothing])
discretization = BayesianPINN([chainl], StochasticTraining(200), param_estim = true,
dataset = [dataset, nothing])

ahmc_bayesian_pinn_pde(pde_system,
discretization;
Expand All @@ -53,9 +47,8 @@ ahmc_bayesian_pinn_pde(pde_system,
saveats = [1 / 50.0],
param = [LogNormal(6.0, 0.5)])

discretization = NeuralPDE.BayesianPINN([chainl],
QuasiRandomTraining(200),
param_estim = true, dataset = [dataset, nothing])
discretization = BayesianPINN([chainl], QuasiRandomTraining(200), param_estim = true,
dataset = [dataset, nothing])

ahmc_bayesian_pinn_pde(pde_system,
discretization;
Expand All @@ -66,8 +59,8 @@ ahmc_bayesian_pinn_pde(pde_system,
saveats = [1 / 50.0],
param = [LogNormal(6.0, 0.5)])

discretization = NeuralPDE.BayesianPINN([chainl],
QuadratureTraining(), param_estim = true, dataset = [dataset, nothing])
discretization = BayesianPINN([chainl], QuadratureTraining(), param_estim = true,
dataset = [dataset, nothing])

ahmc_bayesian_pinn_pde(pde_system,
discretization;
Expand All @@ -78,9 +71,8 @@ ahmc_bayesian_pinn_pde(pde_system,
saveats = [1 / 50.0],
param = [LogNormal(6.0, 0.5)])

discretization = NeuralPDE.BayesianPINN([chainl],
GridTraining([0.02]),
param_estim = true, dataset = [dataset, nothing])
discretization = BayesianPINN([chainl], GridTraining([0.02]), param_estim = true,
dataset = [dataset, nothing])

sol1 = ahmc_bayesian_pinn_pde(pde_system,
discretization;
Expand All @@ -91,18 +83,6 @@ sol1 = ahmc_bayesian_pinn_pde(pde_system,
saveats = [1 / 50.0],
param = [LogNormal(6.0, 0.5)])

discretization = NeuralPDE.BayesianPINN([chainf],
GridTraining([0.02]), param_estim = true, dataset = [dataset, nothing])

sol2 = ahmc_bayesian_pinn_pde(pde_system,
discretization;
draw_samples = 1500,
bcstd = [0.03],
phystd = [0.01], l2std = [0.01],
priorsNNw = (0.0, 1.0),
saveats = [1 / 50.0],
param = [LogNormal(6.0, 0.5)])

param = 2 * π
ts = vec(sol1.timepoints[1])
u_real = [analytic_sol_func1(0.0, t) for t in ts]
Expand All @@ -112,14 +92,6 @@ u_predict = pmean(sol1.ensemblesol[1])
@test mean(u_predict .- u_real) < 0.1
@test sol1.estimated_de_params[1]≈param atol=param * 0.3

ts = vec(sol2.timepoints[1])
u_real = [analytic_sol_func1(0.0, t) for t in ts]
u_predict = pmean(sol2.ensemblesol[1])

@test u_predict≈u_real atol=0.5
@test mean(u_predict .- u_real) < 0.1
@test sol2.estimated_de_params[1]≈param atol=param * 0.3

## Example Lorenz System (Parameter Estimation)
println("Example 2, Lorenz System")
@parameters t, σ_
Expand Down Expand Up @@ -160,15 +132,8 @@ us = us .+ ((0.05 .* randn(size(us))) .* us)
ts_ = hcat(sol(ts).t...)[1, :]
dataset = [hcat(us[i, :], ts_) for i in 1:3]

# using Plots, StatsPlots
# plot(hcat(sol.u...)[1, :], hcat(sol.u...)[2, :], hcat(sol.u...)[3, :])
# plot!(dataset[1][:, 1], dataset[2][:, 1], dataset[3][:, 1])
# plot(dataset[1][:, 2:end], dataset[1][:, 1])
# plot!(dataset[2][:, 2:end], dataset[2][:, 1])
# plot!(dataset[3][:, 2:end], dataset[3][:, 1])

discretization = NeuralPDE.BayesianPINN(chain, NeuralPDE.GridTraining([0.01]);
param_estim = true, dataset = [dataset, nothing])
discretization = BayesianPINN(chain, GridTraining([0.01]); param_estim = true,
dataset = [dataset, nothing])

@named pde_system = PDESystem(eqs, bcs, domains,
[t], [x(t), y(t), z(t)], [σ_], defaults = Dict([p => 1.0 for p in [σ_]]))
Expand All @@ -186,10 +151,5 @@ sol1 = ahmc_bayesian_pinn_pde(pde_system,
idealp = 10.0
p_ = sol1.estimated_de_params[1]

# plot(pmean(sol1.ensemblesol[1]), pmean(sol1.ensemblesol[2]), pmean(sol1.ensemblesol[3]))
# plot(sol1.timepoints[1]', pmean(sol1.ensemblesol[1]))
# plot!(sol1.timepoints[2]', pmean(sol1.ensemblesol[2]))
# plot!(sol1.timepoints[3]', pmean(sol1.ensemblesol[3]))

@test sum(abs, pmean(p_) - 10.00) < 0.3 * idealp[1]
# @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2]
# @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2]
Loading

0 comments on commit 0bd8e1f

Please sign in to comment.