Skip to content

Instantly share code, notes, and snippets.

@drozzy
Created January 2, 2021 01:42
Show Gist options
  • Save drozzy/5396cddea57fd21e5cf47edb9f5c42cb to your computer and use it in GitHub Desktop.
Save drozzy/5396cddea57fd21e5cf47edb9f5c42cb to your computer and use it in GitHub Desktop.
using Printf
using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Quadrature, Cubature, Cuba
@parameters t,x
@variables c(..)
@derivatives Dt'~t
@derivatives Dxx''~x
@derivatives Dx'~x
# @derivatives Dy'~y
# @derivatives Dyy''~y
# Parameters
v = 1
R = 0
D = 0.1 # diffusion
t_max = 2.0
x_min = -1.0
x_max = 1.0
# Equations, initial and boundary conditions
eqs = [ Dt(c(t, x)) ~ D * Dxx(c(t,x)) - Dx(c(t,x)) ]
bcs = [ c(t, x_min) ~ c(t, x_max),
c(0, x) ~ cos(π*x)
]
# Space and time domains
domains = [t ∈ IntervalDomain(0.0,t_max),
x ∈ IntervalDomain(x_min,x_max)
]
# Discretization
nx = 32
dx = (x_max-x_min) / (nx - 1)
dt = 0.01
# Neural network
dim = length(domains)
output = length(eqs)
hidden = 8
chain = FastChain( FastDense(dim, hidden, Flux.σ),
FastDense(hidden, hidden, Flux.σ),
FastDense(hidden, 1))
strategy = GridTraining(dx=dx)
discretization = PhysicsInformedNN(chain, strategy=strategy)
pde_system = PDESystem(eqs, bcs, domains, [t,x], [c])
prob = discretize(pde_system,discretization)
cb = function (p,l)
println("Current loss is: $l")
return false
end
res = GalacticOptim.solve(prob,Optim.BFGS();cb=cb,maxiters=10)
# Plots
phi = discretization.phi
# initθ = discretization.initθ
# acum = [0;accumulate(+, length.(initθ))]
# sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
# minimizers = [res.minimizer[s] for s in sep]
# ts,xs = [domain.domain.lower:dx:domain.domain.upper for domain in domains]
anim = @animate for (i, t) in enumerate(0:dt:t_max)
@info "Animating frame $i..."
c_predict = reshape([phi([t, x], res.minimizer)[1] for x in xs], length(xs))
title = @sprintf("Advection-diffusion t = %.3f", t)
plot(xs, c_predict, label="", title=title) #, ylims=(-1, 1))
end
gif(anim, "advection_diffusion_pinn.gif", fps=15)
c_predict = reshape([ phi([0, x], res.minimizer)[1] for x in xs], length(xs))
plot(xs,c_predict)
c_predict = reshape([ phi([t_max,x], res.minimizer)[1] for x in xs], length(xs))
using Plots
plot(xs,c_predict)
@KirillZubov
Copy link

using Printf
using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Quadrature, Cubature, Cuba

@parameters t,x
@variables c(..)
@derivatives Dt'~t
@derivatives Dxx''~x
@derivatives Dx'~x

# @derivatives Dy'~y
# @derivatives Dyy''~y

# Parameters

v = 1
R = 0
D = 0.1 # diffusion
t_max = 2.0
x_min = -1.0
x_max = 1.0

# Equations, initial and boundary conditions
eqs = [ Dt(c(t, x)) ~ D * Dxx(c(t,x)) - Dx(c(t,x)) ]

bcs = [c(0, x) ~ cos*x),
       c(t, x_min) ~ c(t, x_max)
]

# Space and time domains
domains = [t  IntervalDomain(0.0,t_max),
           x  IntervalDomain(x_min,x_max)
]

dx_= 0.2
# Neural network
dim = length(domains)
output = length(eqs)
hidden = 12

chain = FastChain( FastDense(dim, hidden, Flux.σ),
                    FastDense(hidden, hidden, Flux.σ),
                    FastDense(hidden, 1))

strategy = GridTraining(dx=dx_)

discretization = PhysicsInformedNN(chain, strategy=strategy)

pde_system = PDESystem(eqs, bcs, domains, [t,x], [c])
prob = discretize(pde_system,discretization)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

res = GalacticOptim.solve(prob,Optim.BFGS();cb=cb,maxiters=1000)


# Plots

phi = discretization.phi

# initθ = discretization.initθ

# acum =  [0;accumulate(+, length.(initθ))]
# sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
# minimizers = [res.minimizer[s] for s in sep]
# Discretization
nx = 32
dx = (x_max-x_min) / (nx - 1)
dt = 0.01
ts,xs = [domain.domain.lower:dx:domain.domain.upper for domain in domains]
using Plots
anim = @animate for (i, t) in enumerate(0:dt:t_max)
    @info "Animating frame $i..."
    c_predict = reshape([phi([t, x], res.minimizer)[1] for x in xs], length(xs))
    title = @sprintf("Advection-diffusion t = %.3f", t)
    plot(xs, c_predict, label="", title=title) #, ylims=(-1, 1))
end

gif(anim, "advection_diffusion_pinn.gif", fps=15)

c_predict = reshape([ phi([0, x], res.minimizer)[1] for x in xs], length(xs))
plot(xs,c_predict)

c_predict = reshape([ phi([t_max,x], res.minimizer)[1] for x in xs], length(xs))


plot(xs,c_predict)

advection_diffusion_pinn

@KirillZubov
Copy link

hi @drozzy, it just needs more iteration and less grid

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment