--- a/lattice2d.jl	Thu Jan 06 20:09:14 2022 +0100
+++ b/lattice2d.jl	Thu Jan 06 21:13:22 2022 +0100
@@ -71,6 +71,7 @@
 
         # TO DO: Introduce non-linear coupling
 
+        # Next-neighbor linear coupling
         A[2N+x, 2r-1] += k
         A[2N+x, 2l-1] += k
         A[2N+x, 2i-1] -= 2k
@@ -79,18 +80,29 @@
         A[2N+y, 2b] += k
         A[2N+y, 2i] -= 2k
 
+        # Next-neighbor out-dimensional coupling: i.e., top neighbor causes x displacement.
+        # The coupling constant is just half here. In reality, there would be a nonlinear dependency.
+        A[2N+x, 2t-1] += k/4
+        A[2N+x, 2b-1] += k/4
+        A[2N+x, 2i-1] -= k/2
+
+        A[2N+y, 2r] += k/4
+        A[2N+y, 2l] += k/4
+        A[2N+y, 2i] -= k/2
+
         # Add diagonal neighbors
-        for e in (r1, r2, l1, l2)
-            # Diagonal dependence on √2 neighbors.
-            A[2N+x, 2*e-1] += k/4
-            A[2N+x, 2*e]   += k/8 # Interaction between diag. neighbor y and this element's x is weaker.
-            A[2N+y, 2*e-1] += k/8
-            A[2N+y, 2*e]   += k/4
-        end
-        A[2N+x, x] -= 3/2 * k
-        A[2N+y, y] -= 3/2 * k
+        #
+        # for e in (r1, r2, l1, l2)
+        #     # Diagonal dependence on √2 neighbors.
+        #     A[2N+x, 2*e-1] += k/4
+        #     A[2N+x, 2*e]   += k/8 # Interaction between diag. neighbor y and this element's x is weaker.
+        #     A[2N+y, 2*e-1] += k/8
+        #     A[2N+y, 2*e]   += k/4
+        # end
+        # A[2N+x, x] -= 3/2 * k
+        # A[2N+y, y] -= 3/2 * k
 
-        # damping
+        # Add damping terms.
         A[2N+1:end, 2N+1:end] -= LinearAlgebra.I(2N) * damping
     end
     A
@@ -132,38 +144,46 @@
 
 # These are two different initial configurations:
 #  All elements on two orthogonal sides are displaced initially.
-initial_xy_displacements = vcat(
-    [(1, i, 0., .2) for i in 1:n],
-    [(i, 2, .2, .0) for i in 1:n])
+initial_xy_longitudinal_displacements = vcat(
+    [(1, i, 0., .5) for i in 1:n],
+    [(i, 2, .5, .0) for i in 1:n])
+initial_xy_transversal_displacements = vcat(
+    [(1, i, 0.5, 0.0) for i in 1:n],
+    [(i, 2, 0.0, 0.5) for i in 1:n])
+initial_x_longitudinal_displacements = [(i, 1, .2, 0.) for i in 1:n]
+initial_x_transversal_displacements = [(i, 1, 0., 0.2) for i in 1:n]
 #  One element is displaced diagonally.
 initial_corner_displacements = [(1,1,0.3,0.3), (n, 1, 0.3, -0.3)]
+initial_center_displacements = [(div(n,2), div(n,2), .7, .7)]
 
-u0 = initial_state(n, initial_corner_displacements)
+
+u0 = initial_state(n, initial_xy_transversal_displacements)
 
 # Essential part: the coupling matrix!
 coupling = coupling_matrix(n, 1.; damping=0.0000)
 
 # ODE formalities
-problem = ODEProblem(f2d, u0, (0, 40), coupling)
-@time solution = solve(problem)
+problem = ODEProblem(f2d, u0, (0, 20), coupling);
+@time solution = solve(problem);
 
 @gif for (j, u) in enumerate(solution.u)
     plot_state(n, u, j)
 end
 
-plot(solution, vars=[(0, 1), (0, 2)])
+#plot(solution, vars=[(0, 1), (0, 2)])
 
-# Plot animations for all eigenmodes. Careful: there are n^2 eigenmodes!
-function plot_all_eigenmodes(n=5)
+# Plot animations for all eigenmodes given an n x n array of oscillators.
+# Careful: there are n^2 eigenmodes! (can be restricted to at most `max` eigenmodes being calculated)
+function plot_all_eigenmodes(n=5; max=-1)
     # Enumerate all eigenmodes
     coupling = coupling_matrix(n, 1.; full=true)
     ED = LinearAlgebra.eigen(coupling[2n^2+1:end, 1:2n^2])
     @show size(coupling)
 
-    for i in 1:n^2
+    for i in 1:(max < 0 ? n^2 : max)
         @show i
         ev = abs(ED.values[i])
-        u0 = vcat(ED.vectors[:,i] ./ 3, zeros(2n^2))
+        u0 = vcat(ED.vectors[:,i] ./ 2, zeros(2n^2))
 
         problem = ODEProblem(f2d, u0, (0, 4*2pi/sqrt(ev)), coupling)
         @time sol = solve(problem);
@@ -173,4 +193,25 @@
         end
         gif(a, "anim_$(n)_ev_$(i).gif", fps=12)
     end
+end
+
+# Given a solution, extracts a trajectory for every nth oscillator. Returns two x, y matrices
+# suitable for plotting (each column is one timeseries).
+function extract_2d_trajectories(solution, nth::Int)::Tuple{Matrix,Matrix}
+    N = div(length(solution.u[1]), 4)
+    n = round(Int, sqrt(N))
+    m = div(N, nth)
+    T = length(solution.t)
+    x = zeros(T, m)
+    y = zeros(T, m)
+
+    j = 1
+    for i in 1:nth:N
+        x0, y0 = position(n, i, 1.)
+        x[:,j] = [solution.u[k][2i-1] for k in 1:T] .+ x0
+        y[:,j] = [solution.u[k][2i]   for k in 1:T] .+ y0
+        j += 1
+    end
+
+    (x, y)
 end
\ No newline at end of file