Mercurial > lbo > hg > lattice2d
view linear_chain.jl @ 4:28fbca774ef5
Use sparse arrays, improve plotting for lattice2d
author | Lewin Bormann <lbo@spheniscida.de> |
---|---|
date | Wed, 05 Jan 2022 21:18:21 +0100 |
parents | d7f935c0da5e |
children | 8317f98d7ae6 |
line wrap: on
line source
import LinearAlgebra using DifferentialEquations using Plots theme(:juno) function n_coupled_masses(n::Int, ks::Vector, masses::Vector) # Prepares coupling matrix. if length(ks) != (n+1) @show (length(ks), n+1) error("Length of ks must be n+1.") end if length(masses) != n @show n error("Length of masses must be n.") end A = zeros(4n, 4n) A[1:2n, 2n+1:4n] = LinearAlgebra.I(2n) for i = 1:n # oscillators k1, k2 = ks[i], ks[i+1] x = 2i-1 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] # 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] A[2n+x, x] -= (k1+k2)/masses[i] A[2n+y, y] -= (k1+k2)/masses[i] end A end # The state vector has 4n elements: 2n elements containing x/y positions (alternating, 2i-1 is x), # 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} u0 = zeros(4n) for (di, dx, dy) in displacements u0[2di-1] += dx u0[2di] += dy end u0 end # The differential equation. function f(du::Vector{Float64}, u::Vector{Float64}, A, t) # u has N positions, N velocities. # A gives: du/dt = A u du .= A * u end # Find eigenmodes of system. function eigenmodes(n::Int, coupling::Matrix) sub = @view coupling[2n+1:end, 1:2n] LinearAlgebra.eigen(sub) end # Convert an eigenvector E.vectors[:,i] into a system state. function initial_state_from_eigenmode(em::Vector)::Vector vcat(em, zeros(size(em))) end # Animate a solution of the system. function plot_chain_animation(n::Int, solution) xs = collect(1:n) ys = zeros(n) anim = @animate for (i,u) in enumerate(solution.u) pos = reshape(u[1:2n], 2, n) scatter(xs .+ pos[1,:], pos[2,:], xlim=(-1, N+1), ylim=(-1, 1)) end gif(anim, "anim_$(n).gif") end N = 10 k1, k2 = 1, 2 mass = 1 initial_displacement = [(1, .1, .1), (N, -.1, .1)] timespan = (0, 40) coupling_H = n_coupled_masses(N, [k1, k2, k1, k2, k1, k2, k1, k2, k1, k2, k1], mass * ones(N)) 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) 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") # phase image plot(solution, vars=[(1,2), (3,4), (5,6), (7,8), (9,10)]) savefig("phasefig_$(N).svg") # Animation plot_chain_animation(N, solution)