Is there an example of solving the Schrodinger equation?

[Runfa-Zhang]

like Schrodinger example as follows Fig,

[ChrisRackauckas]

It would be good to add.

[KirillZubov]

I tried to write this a long time ago, here is a draft(not work and overdate)

## Shrödinger Equation

##Let h = u + iv where \[u,v\] = phi(x,t), then equation:
#can be represent to the system:

@parameters t x θ
@variables u(..),v(..)
@derivatives Dt'~t
@derivatives Dx'~x
@derivatives Dxx''~x

# eq = 1im*Dt(h(t,x,θ)) + 0.5*Dxx(h(t,x,θ)) + abs2(h(t,x,θ))*h(t,x,θ) ~ 0.
#let h = u + iv where [u,v] = phi(x,t), then eq can be represent to the system
eqs =[Dt(u(t,x,θ)) + 0.5*Dxx(v(t,x,θ)) + (u(t,x,θ)^2 + v(t,x,θ)^2)*v(t,x,θ) ~ 0.,
      Dt(v(t,x,θ)) - 0.5*Dxx(u(t,x,θ)) - (u(t,x,θ)^2 + v(t,x,θ)^2)*u(t,x,θ) ~ 0.]

# Initial and boundary conditions
bcs = [u(0,x,θ) ~ 2*sech(x),v(0,x,θ) ~ 0.
       u(t,-5,θ) ~ 0.,v(t,-5,θ) ~ 0.,
       u(t,5,θ) ~ 0.,v(t,5,θ) ~ 0.
       Dx(u(t,-5,θ)) ~ 0., Dx(v(t,-5,θ)) ~ 0.
       Dx(u(t,5,θ)) ~ 0., Dx(v(t,5,θ)) ~ 0.
       ]
# Space and time domains
domains = [t ∈ IntervalDomain(0.0,pi/2),
           x ∈ IntervalDomain(-5.,5.)]
# Discretization
dx = 0.2;dt = 0.1
# Neural network

chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,2))


strategy = NeuralPDE.GridTraining()
discretization = NeuralPDE.PhysicsInformedNN([dt,dx],
                                             chain,
                                             strategy= strategy)

pde_system = PDESystem(eqs,bcs,domains,[t,x],[u,v])
prob = NeuralPDE.discretize(pde_system,discretization)

# res = GalacticOptim.solve(prob, ADAM(0.1), progress = false; cb = cb, maxiters=10000)
res = GalacticOptim.solve(prob,BFGS(); cb = cb, maxiters=2000)
phi = discretization.phi

ts,xs = [domain.domain.lower:dx:domain.domain.upper for (dx, domain) in zip([dt,dx],domains)]

h_abs(t,x) = sqrt(phi([t,x],res.minimizer)[1]^2+phi([t,x],res.minimizer)[2]^2)
h_predict = reshape([h_abs(t,x) for t in ts for x in xs],length(xs),length(ts))
h_predict2 = [[h_abs(t,x)  for x in xs] for t in ts]
plot(ts, xs, h_predict,  linetype=:contourf)
p1=plot(xs, h_predict2[1]);
p2=plot(xs, h_predict2[4]);
p3=plot(xs, h_predict2[end-1]);
plot(p1,p2,p3)

[Runfa-Zhang]

I tried to write this a long time ago, here is a draft(not work and overdate)

## Shrödinger Equation

##Let h = u + iv where \[u,v\] = phi(x,t), then equation:
#can be represent to the system:

@parameters t x θ
@variables u(..),v(..)
@derivatives Dt'~t
@derivatives Dx'~x
@derivatives Dxx''~x

# eq = 1im*Dt(h(t,x,θ)) + 0.5*Dxx(h(t,x,θ)) + abs2(h(t,x,θ))*h(t,x,θ) ~ 0.
#let h = u + iv where [u,v] = phi(x,t), then eq can be represent to the system
eqs =[Dt(u(t,x,θ)) + 0.5*Dxx(v(t,x,θ)) + (u(t,x,θ)^2 + v(t,x,θ)^2)*v(t,x,θ) ~ 0.,
      Dt(v(t,x,θ)) - 0.5*Dxx(u(t,x,θ)) - (u(t,x,θ)^2 + v(t,x,θ)^2)*u(t,x,θ) ~ 0.]

# Initial and boundary conditions
bcs = [u(0,x,θ) ~ 2*sech(x),v(0,x,θ) ~ 0.
       u(t,-5,θ) ~ 0.,v(t,-5,θ) ~ 0.,
       u(t,5,θ) ~ 0.,v(t,5,θ) ~ 0.
       Dx(u(t,-5,θ)) ~ 0., Dx(v(t,-5,θ)) ~ 0.
       Dx(u(t,5,θ)) ~ 0., Dx(v(t,5,θ)) ~ 0.
       ]
# Space and time domains
domains = [t ∈ IntervalDomain(0.0,pi/2),
           x ∈ IntervalDomain(-5.,5.)]
# Discretization
dx = 0.2;dt = 0.1
# Neural network

chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,2))


strategy = NeuralPDE.GridTraining()
discretization = NeuralPDE.PhysicsInformedNN([dt,dx],
                                             chain,
                                             strategy= strategy)

pde_system = PDESystem(eqs,bcs,domains,[t,x],[u,v])
prob = NeuralPDE.discretize(pde_system,discretization)

# res = GalacticOptim.solve(prob, ADAM(0.1), progress = false; cb = cb, maxiters=10000)
res = GalacticOptim.solve(prob,BFGS(); cb = cb, maxiters=2000)
phi = discretization.phi

ts,xs = [domain.domain.lower:dx:domain.domain.upper for (dx, domain) in zip([dt,dx],domains)]

h_abs(t,x) = sqrt(phi([t,x],res.minimizer)[1]^2+phi([t,x],res.minimizer)[2]^2)
h_predict = reshape([h_abs(t,x) for t in ts for x in xs],length(xs),length(ts))
h_predict2 = [[h_abs(t,x)  for x in xs] for t in ts]
plot(ts, xs, h_predict,  linetype=:contourf)
p1=plot(xs, h_predict2[1]);
p2=plot(xs, h_predict2[4]);
p3=plot(xs, h_predict2[end-1]);
plot(p1,p2,p3)

Thanks very much！

[munozariasjm]

I have also worked with it. Much of the cases for Schrodinger equation it is important to keep an eye on the normalization. I have a little demo here for the TISE. Hope it works for you

using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim
using DiffEqFlux, Plots, Statistics, DomainSets, Quadrature
import ModelingToolkit: Interval, infimum, supremum

function sol_Schr_1D(𝔼,nits)

    #Geometrical configuration
    x_0 = 0.
    x_end = 1.
    # symbolic parameters
    @parameters x
    @variables u(..)
    Dxx = Differential(x)^2
    Dx = Differential(x)
    𝕍(x) = (x-0.5)^2
    eq  = Dxx(u(x)) ~ -2 .*( 𝔼 .- 𝕍(x))*u(x)
    rand_num = rand()
    # bcs = [u(0.) ~ 0.0,u(1.) ~ 0.0, Dx(u(0.5))~0.]

    lb = [x_0]
    ub = [x_end]
    function norm_loss_function(phi,θ,p)
        function inner_f(x,θ)
             dx*(phi(x, θ).^2) .- 1
        end
        prob = QuadratureProblem(inner_f, lb, ub, θ)
        norm2 = solve(prob, HCubatureJL(), reltol = 1e-8, abstol = 1e-8, maxiters =10);
        abs(norm2[1])
        # f(x) = phi(x)^2
        # prob = QuadratureProblem(f,lb,ub)
        # sol = solve(prob,HCubatureJL(),reltol=1e-3,abstol=1e-3)
    end

    bcs = [u(0.) ~ 0.0,u(1.) ~ 0.0,Dx(u(0.5))~ 0]

    domains = [x ∈ Interval(0.0,1.0)]
    dim = 1 # number of dimensions
    chain = FastChain(FastDense(dim,8,Flux.σ),
                      FastDense(8,16,Flux.σ),
                      FastDense(16,16,Flux.σ),
                      FastDense(16,1))

    initθ = Float64.(DiffEqFlux.initial_params(chain))
    eltypeθ = eltype(initθ)
    dx = 0.1
    strategy = GridTraining(dx)#QuasiRandomTraining(20)

    discretization = PhysicsInformedNN(chain, strategy;
                                        init_params = initθ,
                                        additional_loss=norm_loss_function)

    @named pde_system = PDESystem(eq,bcs,domains,[x],[u(x)])
    prob = discretize(pde_system,discretization)

    pde_inner_loss_functions = prob.f.f.loss_function.pde_loss_function.pde_loss_functions.contents
    bcs_inner_loss_functions = prob.f.f.loss_function.bcs_loss_function.bc_loss_functions.contents

    print(("⇊"^50)*("\n"))
    print("     Now let me think... tryinig to learn       \n")
    print(("⇈"^50)*("\n"))
    phi = discretization.phi
    function cb_(p,l)
        println("loss: ", l )
        println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
        println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
        return false
    end

    # res = GalacticOptim.solve(prob,LBFGS(),cb = cb_,maxiters=400)
    opt = Optim.BFGS()
    # opt = ADAM(0.05)
    # prob = remake(prob,u0=res.minimizer)

    res = GalacticOptim.solve(prob, opt, cb = cb_, maxiters=Int(nits))

    indvars = [x]
    depvars = [u]
    xs = [infimum(d.domain):dx/10:supremum(d.domain) for d in domains][1]
    # xs = Vector(-1.: dx:1+dx)

    u_predict  = [first(phi(x,res.minimizer)) for x in xs]

    u_predict
end

ψ = sol_Schr_1D(1., 1e5)
plot((Vector(0.:1/(100):1.),ψ.^2),label="ψ")
plot!((0.:1/(200):1.),(Vector(0.:1/(200):1.).- 0.5).^2,label="Potential")