Untitled

module fDNLS_Direct
using DifferentialEquations, GLMakie, LinearAlgebra, OffsetArrays

using SpecialFunctions: zeta
GLMakie.activate!()

#=
The fDNLS is written as iuₙ' = ∑(uₙ - uₘ)/|n-m|^(1+α) - |uₙ|^2 uₙ, where the sum is evaluated such that -N <= m <= N and m != n. 
Dirichlet boundary values may be imposed: u_(-N) = 0, u_(N) = 0.

Standing wave solutions are of the form u(x,t) = q(x)exp[i ω t], where q(x) here is the "Ground State Solution" obtained by solution by the 
corresponding system of nonlinear equations after substituting the ansatz. To solve the above system of ODEs, take u(x,0) = q(x). Ref. https://royalsocietypublishing.org/doi/epdf/10.1098/rspa.2014.0364
=#

#=
DifferentialEquations - Used for the ODE Solver
GLMakie               - 2D/3D plotting backend,
LinearAlgebra,
SpecialFunctions      - Used for importing the Riemann zeta function
=#

function main(N)
    #=
    Per the boundary conditions, solve for u = (u_{-N+1}, ..., u_{N-1})
    =#

    h      = 1
    α      = 0.2 
    tspan = (0.0,25.0)
    x = -N+1:N-1

    # Per u_(-N) = 0 and u_(N) = 0, then we solve the (2N+1) - 2 = 2N-1 total equations
    RHS(u,p,t) = im*(h*coefficientMatrix(α,x)*u - abs2.(u).*u)

   # initial_values = 0.9*Complex.(sech.(x)).* exp.(im*x)

    initial_values = onSite(30,α,x)# .* exp.(im*x) # onSite(ω, α)
    # Equivalent to Ode45 from Matlab
    prob = ODEProblem(RHS, initial_values, tspan)
    sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8)   

    plot_main(sol, x) # Surface and contour plot of |u|²
    return sol
end

function coefficientMatrix(α,x) 

    M = OffsetArray(zeros(length(x), length(x)), x,x)
    
    for i in firstindex(M,1) : lastindex(M,1)
        for j in axes(M,2)
            i == j ? M[i,i] = 1/abs(i - (-N))^(1+α) + sum(1/abs(i - m)^(1+α) for m in x if m!=i) + 1/abs(i - N) : M[i,j] = -1/abs(i - j)
        end
    end

    return M.parent   
end

function onSite(ω,α,x)
    # Assumption: qₙ = q₋ₙ, q₀ >> 1 >> q₁ >> ...
    Q = OffsetVector(zeros(Complex{Float64},length(x)),x)
    Q[0] = sqrt(ω + 2*zeta(1+α))

    # Obtain q₂, q₃, ... via equation (3.1)
    for n in 1:maximum(x) #drop first 1 since already initialized
        asymptotic_onsite = sum(Q[j]/(n+j)^(1+α) for j in 1-n:n-1)/(2*zeta(1+α) - 1/(2n)^(1+α) + ω)
        Q[-n] = asymptotic_onsite
        Q[n]  = asymptotic_onsite
    end

    return Q.parent # typeof(Q) = OffsetVector, typeof(Q.parent) = Vector
end

function plot_main(sol,x)
    fig = Figure(backgroundcolor=:snow2)

    y = sol.t
    z= [abs(sol[i,j])^2 for i in axes(sol,1), j in axes(sol,2)]  # sol[i,j] is the ith component at timestep j.

    p1 = surface(fig[1,1], x, y, z,
        axis=(type=Axis3, xlabel="Space", ylabel="Time", zlabel="|u|²", 
            title=L"\text{Intensity plot with } u(x,0) = \text{ asymptotic onsite sequence}"))

    #p2 =  contourf(fig[1,2], x, y, z, axis=(xlabel="Space", ylabel="Time"))
    p2  =  lines(fig[1,2], x, onSite(30,0.2,x), xlabel="Space", ylabel="qₙ", title="Asymptotic onsite sequence")

    normSol_inf = [norm(sol[i], Inf) for i in 1:length(sol.t)]

    p3 = lines(fig[2,:], y, normSol_inf, label=L"k=1", axis=(xlabel="Time", ylabel=L"L^{\infty}", xticks=0:0.5:20))
    axislegend()
    display(fig)
end

const N=9
main(N)
end