--- a/linear_chain.jl	Thu Jan 06 21:13:22 2022 +0100
+++ b/linear_chain.jl	Sun Jan 09 11:11:05 2022 +0100
@@ -6,7 +6,7 @@
 
 # The coupling matrix has 4nx4n elements for a linear chain of two degrees of freedom per oscillator.
 # We use Born-von Karman boundary conditions, i.e. the chain is a ring.
-function n_coupled_masses(n::Int, ks::Vector, masses::Vector)
+function n_coupled_masses(n::Int, ks::Vector, masses::Vector; fixed_ends=false)
     # Prepares coupling matrix.
     if length(ks) != (n+1)
         @show (length(ks), n+1)
@@ -27,12 +27,25 @@
         y = 2i
 
         # Link previous oscillator
-        A[2n+x, (x-2) < 1 ? (2n-1) : (x-2)] += k1/masses[i]
-        A[2n+y, (y-2) < 1 ? (2n) : (y-2)] += k1/masses[i]
+        prevx = (x-2) < 1 ? (2n-1) : (x-2)
+        prevy = (y-2) < 1 ? (2n) : (y-2)
+        if !fixed_ends || ((x-2) >= 1)
+            A[2n+x, prevx] += k1/masses[i]
+        end
+        if !fixed_ends || ((y-2) >= 1)
+            A[2n+y, prevy] += k1/masses[i]
+        end
+
 
         # Link next oscillator
-        A[2n+x, ((x+2)%(2n))] += k2/masses[i]
-        A[2n+y, max(((y+2-1)%(2n)+1), 2)] += k2/masses[i]
+        nextx = ((x+2)%(2n))
+        nexty = max(((y+2-1)%(2n)+1), 2)
+        if !fixed_ends || (nextx == x+2)
+            A[2n+x, nextx] += k2/masses[i]
+        end
+        if !fixed_ends || (nexty == y+2)
+            A[2n+y, nexty] += k2/masses[i]
+        end
 
         A[2n+x, x] -= (k1+k2)/masses[i]
         A[2n+y, y] -= (k1+k2)/masses[i]
@@ -44,7 +57,7 @@
 # and another 2n elements containing x/y velocities (alternating).
 # The initial displacements are given as (i, x, y) tuples representing a
 # (x,y) displacement of the i'th oscillator.
-function initial_state(n::Int, displacements::Vector{Tuple{Int,Float64,Float64}})::Vector{Float64}
+function initial_state(n::Int, displacements::Vector)::Vector{Float64}
     u0 = zeros(4n)
     for (di, dx, dy) in displacements
         u0[2di-1] += dx
@@ -84,32 +97,41 @@
     gif(anim, "anim_$(n).gif")
 end
 
-N = 10
-k1, k2 = 1, 2
+function initial_wave_displacement(N::Int, q::Float64, ampl::Float64=.25, a::Float64=1.; trans=false)::Vector{Tuple}
+    N = div(N, 2)
+    ampls = [ampl * sin(q * a * i) for i in 1:N]
+    [(i, trans ? 0. : ampls[i], trans ? ampls[i] : 0.) for i in 1:N]
+end
+
+N = 15
+k1, k2 = .4, 1
 mass = 1
-initial_displacement = [(1, .4, .0), (N, -.0, .0)]
-timespan = (0, 20)
-coupling_H = n_coupled_masses(N, [k1, k2, k1, k2, k1, k2, k1, k2, k1, k2, k1], mass * ones(N))
+initial_displacement = [(1, .0, .3), (N, -.0, .0)]
+initial_displacement_2 = [(1, .5, 0.)]
+initial_wave_4 = initial_wave_displacement(N, 2pi/(N), .1; trans=true)
+timespan = (0, 90)
+coupling_H = n_coupled_masses(N, k1 * ones(N+1), mass * ones(N); fixed_ends=false)
 
 E = eigenmodes(N, coupling_H)
 
 # Two different initial configurations as choice.
 u0 = initial_state(N, initial_displacement)
-u1 = initial_state_from_eigenmode(E.vectors[:,14] ./ 3)
+u1 = initial_state_from_eigenmode(E.vectors[:,1] ./ 3)
+u2 = initial_state(N, initial_wave_4)
 
 problem = ODEProblem(f, u0, timespan, coupling_H)
 @time solution = solve(problem)
 
 # 3D trajectory plot
-plot(solution, vars=[(0, 2x-1, 2x) for x in 1:N],
-     size=(900,450), title="x trajectories")
-savefig("trajectory_3d_$(N).svg")
+# plot(solution, vars=[(0, 2x-1, 2x) for x in 1:N],
+    #  size=(900,450), title="x trajectories")
+# savefig("trajectory_3d_$(N).svg")
 
 # phase image
-plot(solution, vars=[(1,2), (3,4), (5,6), (7,8), (9,10)])
-savefig("phasefig_$(N).svg")
+# plot(solution, vars=[(1,2), (3,4), (5,6), (7,8), (9,10)])
+# savefig("phasefig_$(N).svg")
 
-plot(solution, vars=[(0, 2i-1) for i in 1:N])
+# plot(solution, vars=[(0, 2i-1) for i in 1:N])
 
 # Animation
 plot_chain_animation(N, solution)