Skip to content

Commit

Permalink
Merge pull request #13 from oashour/type_fix
Browse files Browse the repository at this point in the history
Fix types to use parametric types
  • Loading branch information
oashour committed Oct 20, 2020
2 parents a096670 + d6c3b86 commit bf5acea
Show file tree
Hide file tree
Showing 8 changed files with 271 additions and 253 deletions.
1 change: 1 addition & 0 deletions src/NLSS.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
module NLSS

include("Simulation.jl")
include("Utilities.jl")
include("Plotter.jl")

end #module
26 changes: 13 additions & 13 deletions src/Plotter.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,33 +7,33 @@ using FFTW
export plot_ψ, plot_ψ̃, plot_CoM

"""
function plot_CoM(sim::Sim, x_res::Int64 = 500)
function plot_IoM(sim::Sim, x_res::Int64 = 500)
Plots the integrals of motion for a `Sim` object `sim` with a resolution `x_res` points in
the x-direction. Produces plots of the energy, kinetic energy, potential energy, energy
error, particle number and momentum
See also: [`Simulation.compute_CoM!`](@ref)
"""
function plot_CoM(sim::Sim, x_res = 500)
function plot_IoM(obj, x_res = 500)
println("Plotting energy with a resolution of $x_res")
xₛ = Int(ceil(sim.box.Nₜ/x_res))
p1 = plot(sim.box.x[1:xₛ:end], [sim.E[1:xₛ:end], sim.KE[1:xₛ:end], sim.PE[1:xₛ:end]], label = [L"E" L"T" L"V"])
xₛ = Int(ceil(obj.box.Nₜ/x_res))
p1 = plot(obj.box.x[1:xₛ:end], [obj.E[1:xₛ:end], obj.KE[1:xₛ:end], obj.PE[1:xₛ:end]], label = [L"E" L"T" L"V"])
xlabel!(L"x")
ylabel!(L"E")
xlims!((minimum(sim.box.x), maximum(sim.box.x)))
p2 = plot(sim.box.x[1:xₛ:end], sim.dE[1:xₛ:end], label = "")
xlims!((minimum(obj.box.x), maximum(obj.box.x)))
p2 = plot(obj.box.x[1:xₛ:end], obj.dE[1:xₛ:end], label = "")
xlabel!(L"x")
ylabel!(L"\delta E")
xlims!((minimum(sim.box.x), maximum(sim.box.x)))
p3 = plot(sim.box.x[1:xₛ:end], [sim.P[1:xₛ:end]], label = "")
xlims!((minimum(obj.box.x), maximum(obj.box.x)))
p3 = plot(obj.box.x[1:xₛ:end], [obj.P[1:xₛ:end]], label = "")
xlabel!(L"x")
ylabel!(L"P")
xlims!((minimum(sim.box.x), maximum(sim.box.x)))
p4 = plot(sim.box.x[1:xₛ:end], [sim.N[1:xₛ:end]], label = "")
xlims!((minimum(obj.box.x), maximum(obj.box.x)))
p4 = plot(obj.box.x[1:xₛ:end], [obj.N[1:xₛ:end]], label = "")
xlabel!(L"x")
ylabel!(L"N")
xlims!((minimum(sim.box.x), maximum(sim.box.x)))
xlims!((minimum(obj.box.x), maximum(obj.box.x)))

return p1, p2, p3, p4
end
Expand All @@ -49,7 +49,7 @@ The function can plot either a heatmap (`mode = "density"`) or a 3D surface
See also: [`Simulation.solve!`](@ref)
"""
function plot_ψ(sim::Sim; mode = "density", power=1, x_res=500, t_res=512)
function plot_ψ(sim; mode = "density", power=1, x_res=500, t_res=512)
if ~sim.solved
throw(ArgumentError("The simulation has not been solved, unable to plot."))
end
Expand Down Expand Up @@ -100,7 +100,7 @@ line if `mode = lines`. The latter uses `x_res` points as well.
See also: [`Simulation.compute_spectrum!`](@ref)
"""
function plot_ψ̃(sim::Sim; mode = "density", x_res=500, ω_res=512, skip = 1, n_lines = 10)
function plot_ψ̃(sim; mode = "density", x_res=500, ω_res=512, skip = 1, n_lines = 10)
if ~sim.solved
throw(ArgumentError("The simulation has not been solved, unable to plot."))
end
Expand Down
196 changes: 196 additions & 0 deletions src/SimExtras.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,196 @@
export Sim, Box
export params, print
export ψ₀_periodic

###########################################################################
# Types
###########################################################################
struct Box{TT<:Real}
t::Array{TT, 1}
ω::Array{TT, 1}
x::Array{TT, 1}
Nₜ::Int64
Nₓ::Int64
dt::TT
dx::TT
n_periods::Int64
end #SimulationBox

function Box(xᵣ::Pair, T; dx = 1e-3, Nₜ = 256, n_periods = 1)
println("==========================================")
println("Initializing simulation box with $n_periods period(s) and dx = $dx, Nₜ = $Nₜ.")
T = n_periods * T
println("Longitudinal range is [$(xᵣ.first), $(xᵣ.second)], transverse range is [$(-T/2), $(T/2))")
dt = T / Nₜ
t = dt * collect((-Nₜ/2:Nₜ/2-1))

x = collect(xᵣ.first:dx:xᵣ.second)
Nₓ = length(x)
ω = 2π/T * collect((-Nₜ/2:Nₜ/2-1))

box = Box(t, ω, x, Nₜ, Nₓ, dt, dx, n_periods)

println("Done computing t, x, ω")
println("==========================================")

return box
end

mutable struct Sim{TT<:Real}
λ::Complex{TT}
T::TT
Ω::TT
box::Box{TT}
ψ₀::Array{Complex{TT}, 1}
algorithm::String
αₚ::TT
solved::Bool
ψ::Array{Complex{TT}, 2}
spectrum_computed::Bool
ψ̃::Array{Complex{TT}, 2}
energy_computed::Bool
E::Array{TT, 1}
PE::Array{TT, 1}
KE::Array{TT, 1}
dE::Array{TT, 1}
N::Array{TT, 1}
P::Array{TT, 1}
end # Simulation

function Sim(λ, box::Box, ψ₀::Array{Complex{TT}, 1}; algorithm = "2S", αₚ = 0.0) where TT <: Real
ψ = Array{Complex{TT}}(undef, box.Nₓ, box.Nₜ)
ψ̃ = similar(ψ)
E = zeros(box.Nₓ)
PE = similar(E)
KE = similar(E)
dE = similar(E)
N = similar(E)
P = similar(E)
if αₚ < 0.0
throw(ArgumentError("αₚ < 0. Set αₚ = 0 to disable pruning or αₚ = Inf for fixed pruning. See documentation)"))
end
if αₚ > 0.0 && box.n_periods == 1
throw(ArgumentError("Pruning is only applicable when n_periods > 1."))
end
# Compute some parameters
λ, T, Ω = params= λ)
sim = Sim(λ, T, Ω, box, ψ₀, algorithm, αₚ, false, ψ, false, ψ̃, false, E, PE, KE, dE,
N, P)
return sim
end #init_sim

###########################################################################
# Utility
###########################################################################
function params(; kwargs...)
if length(kwargs) != 1
throw(ArgumentError("You have either specified too few or too many parameters. You must specify one and only one of the following options: λ, Ω, T, a."))
end
param = Dict(kwargs)
if :a in keys(param)
λ = im * sqrt(2 * param[:a])
T = π/sqrt(1 - imag(λ)^2)
Ω = 2π/T
println("Passed a = $(param[:a]), computed λ = , T = $T and Ω = ")
elseif in keys(param)
λ = param[]
T = π/sqrt(1 - imag(λ)^2)
Ω = 2π/T
println("Passed λ=, computed T = $T and Ω = ")
elseif in keys(param)
λ = im * sqrt((1 - (param[] / 2)^2))
T = π/sqrt(1 - imag(λ)^2)
Ω = param[]
println("Passed Ω=, computed λ = and T = $T")
elseif :T in keys(param)
λ = im * sqrt((1 - ((2*π/param[:T]) / 2)^2))
T = param[:T]
Ω = 2π/T
println("Passed T = $T, computed λ = and Ω = ")
end

return λ, T, Ω
end #compute_parameters

"""
function print(sim::Sim)
Prints information about the `Sim` instance `sim`
"""
function print(sim::Sim)
println("Box Properties:")
println("------------------------------------------")
println("dx = $(sim.box.dx) (Nₓ = $(sim.box.Nₓ))")
println("Nₜ = $(sim.box.dx) (dt = $(sim.box.dt))")
println("Parameters:")
println("------------------------------------------")
println("λ = $(sim.λ)")
println("Ω = $(sim.Ω)")
println("T = $(sim.T)")
println("------------------------------------------")
# Should add information about ψ₀
if sim.algorithm == "2S"
println("Algorithm: second order symplectic")
elseif sim.algorithm == "4S"
println("Algorithm: fourth order symplectic")
elseif sim.algorithm == "6S"
println("Algorithm: sixth order symplectic")
elseif sim.algorithm == "8S"
println("Algorithm: eighth order symplectic")
else
throw(ArgumentError("Algorithm type unknown, please check the documentation"))
end
println("------------------------------------------")
end

###########################################################################
# ψ₀
###########################################################################
"""
function ψ₀_periodic(coeff, box::SimBox, params::SimParamseters; phase=0)
Computes an initial wavefunction for the `SimBox` `box`, with fundamental frequency `sim.Ω`
and coefficients ``A_1...n`` = `coeff` and an overall phase `exp(i phase t)`, i.e. of the form:
**TODO**: insert latex form here
See also: [`init_sim`](@ref)
"""
function ψ₀_periodic(coeff::Array, box::Box, Ω; phase=0)
println("==========================================")
println("Initializing periodic ψ₀")
for (n, An) in enumerate(coeff)
if abs(An) >= 1
println("The absolute value of the coefficient A($(n+1)) = $(An) is greater than 1. psi_0 needs to be normalizable.")
else
println("A($(n)) = $(An)")
end #if
end #for

println("Computing A₀ to preserve normalization.")
A0 = sqrt(1 - 2 * sum([abs(An) .^ 2 for An in coeff]))
println("Computed A₀ = $A0")
#A0 = 1
if phase != 0
str = "ψ₀ = exp(i $phase t) ($A0 + "
else
str = "ψ₀ = $A0 + "
end

ψ₀ = A0 * ones(box.Nₜ)
for (n, An) in enumerate(coeff)
ψ₀ += 2 * An * cos.(n * Ω * box.t)
str = string(str, "2 × $An × cos($n × $(Ω) t)")
end #for
# Multiply by the overall phase
ψ₀ = exp.(im * phase * box.t) .* ψ₀

if phase != 0
str = string(str, ")")
end

println(str)
println("==========================================")

return ψ₀
end #psi0_periodic
5 changes: 4 additions & 1 deletion src/SimSolvers.jl
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
export solve!

"""
solve!(sim::Simulation)
Expand Down Expand Up @@ -55,6 +55,9 @@ function solve!(sim::Sim)
return nothing
end #solve

#######################################################################################
# Algorithms
#######################################################################################
"""
T2(ψ, ω, dx, F)
Expand Down
107 changes: 0 additions & 107 deletions src/SimTypes.jl

This file was deleted.

Loading

0 comments on commit bf5acea

Please sign in to comment.