Simple 2D lattice dynamics calculation for a cubic structure with next-neighbor coupling only. Example: https://borgac.net/~lbo/asset/jl_62ec52.gif. Eigenmmodes: https://borgac.net/~lbo/asset/eigenmodes_5x5/

lattice2d.jl

# (c) 2022 Lewin Bormann
# Full and updated code at https://borgac.net/lbo/hg/lattice2d/

using Plots
using DifferentialEquations
import LinearAlgebra

using SparseArrays

theme(:juno)

# Masses schema:
#
#  7 -- 8 -- 9
#  |  X |  X |
#  4 -- 5 -- 6
#  |  X |  X |
#  1 -- 2 -- 3
# 
# where 'X' is a cross connection.
# We use Born-van Karman boundary conditions, i.e.
# a "Pacman world" where the area behaves like it is on a torus.

function neighbors(n::Int, i::Int)::Tuple
    # Returns neighbors in bottom/top/right/left
    a = (i + n) % n^2
    a = (a==0) ? n^2 : a

    b = (i - n) < 1 ? (n*(n-1)+i) : (i-n)


    c = i + 1 % n^2
    c = (div(c-1, n) > div(i-1, n)) ? (i-n+1) : c

    d = (i - 1) < 1 ? n : (i-1)
    d = (div(d-1, n) < div(i-1, n)) ? (i+n-1) : d

    (a,b,c,d)
end

function position(n::Int, i::Int, a::Float64)::Tuple
    ((i-1)%n, (div(i-1, n)))
end

# Each oscillator has two equations: one for x and one for the y direction.
# Assuming the neighbors L/R (x), T/B (y), we get the following equations of motion for oscillator i:
#
# ̈x_i = k (x_R + x_L - 2 x_i)
# ̈y_i = k (y_T + y_B - 2 y_i)
#
# This assumes constant spring constant k for all springs, obviously.
# In the vector notation, oscillator i has its x component at index 2i-1 and its y component at 2i.
# Additional terms for diagonal springs are introduced. The spring constants for this are lower.

function coupling_matrix(n::Int, k::Float64; damping::Float64=0., full=false)
    N = n^2
    if full
        A = zeros(4N, 4N)
    else
        A = spzeros(4N, 4N)
    end

    # Link velocity to displacement.
    A[1:2N, 2N+1:end] = LinearAlgebra.I(2N)

    for i in 1:N
        (b,t,r,l) = neighbors(n, i)
        # Diagonal neighbors
        (_, _, r1, l1) = neighbors(n, t)
        (_, _, r2, l2) = neighbors(n, b)

        x = 2i-1
        y = 2i

        # TO DO: Introduce non-linear coupling
        # TO DO: Introduce damping

        A[2N+x, 2r-1] += k
        A[2N+x, 2l-1] += k
        A[2N+x, 2i-1] -= 2k

        A[2N+y, 2t] += k
        A[2N+y, 2b] += k
        A[2N+y, 2i] -= 2k

        # 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

        # damping
        A[2N+1:end, 2N+1:end] -= LinearAlgebra.I(2N) * damping
    end
    A
end

function initial_state(n::Int, displacements::Vector{Tuple{Int, Int, Float64, Float64}})::Vector
    u0 = zeros(4n^2)
    N = n^2
    for (x, y, dx, dy) in displacements
        i = n*(x-1) + y
        ix, iy = 2i-1, 2i
        #@show (i, ix, iy)
        u0[ix] += dx
        u0[iy] += dy
    end
    u0
end

function plot_state(n::Int, u::Vector, i::Int=0; a=1.)
    N = n^2
    dis = u[1:2N]
    dxy = reshape(dis, 2, N)

    is = collect(1:N)
    xs = ((is .- 1) .% n)
    ys = div.(is .- 1, n)

    xs += dxy[1,:]
    ys += dxy[2,:]
    scatter(xs .* a, ys .* a, xlim=(-a, n*a), ylim=(-a, n*a), title="Positions $(i)")
end

function f2d(du, u, A, t)
    LinearAlgebra.mul!(du, A, u)
end

# There will be n x n oscillators.
n = 10

# 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])
#  One element is displaced diagonally.
initial_corner_displacements = [(1,1,0.3,0.3), (n, 1, 0.3, -0.3)]

u0 = initial_state(n, initial_corner_displacements)

# Essential part: the coupling matrix!
coupling = coupling_matrix(n, 1.; damping=0.0001)

# ODE formalities
problem = ODEProblem(f2d, u0, (0, 40), 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 animations for all eigenmodes. Careful: there are n^2 eigenmodes!
function plot_all_eigenmodes(n=5)
    # 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
        @show i
        ev = abs(ED.values[i])
        u0 = vcat(ED.vectors[:,i] ./ 3, zeros(2n^2))

        problem = ODEProblem(f2d, u0, (0, 4*2pi/sqrt(ev)), coupling)
        @time sol = solve(problem);

        a = @animate for (j,u) in enumerate(sol.u)
            plot_state(n, u, j)
        end
        gif(a, "anim_$(n)_ev_$(i).gif", fps=12)
    end
end