Animating phase space movement in a 1D harmonic oscillator: https://youtu.be/v4edySA8Gzk (contains numeric friction)

harmonic_oscillator_phasespace_animate.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Dec  1 12:12:22 2020

@author: lbo
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.animation
import math

from collections import namedtuple

class Simulator:
    Config = namedtuple('Config', ('n', 'x0', 'deltax', 'v0', 'deltav', 'omega'))

    def __init__(self, n=20, x0=10, deltax=3, v0=0, deltav=3, omega=1):
        self.cfg = Simulator.Config(n=n, x0=x0, deltax=deltax, v0=v0, deltav=deltav, omega=omega)
        self.pos = x0 + (1/2 - np.random.random(n)) * deltax
        self.v = (1/2 - np.random.random(n)) *deltav + v0

    def increment(self, dt):
        """Returns tuple (positions, velocities)."""
        self.pos += self.v * dt - 1/2 * self.pos * self.cfg.omega**2 * dt**2
        self.v -= dt * self.cfg.omega**2 * self.pos


    def animate(self, dt=0.1, fps=10, n=1000, dst=None):
        fig, ax = plt.subplots()

        xs, vs = self.pos, self.v
        scat = plt.scatter(xs, vs, [0.5]*len(xs), color='blue')
        ax.set_xlim(-self.cfg.x0*1.5, self.cfg.x0*1.5)
        ax.set_ylim(-20, 20)
        ax.grid()

        # Depends on positive initial state
        emax = self.v.max()**2/2 + self.cfg.omega**2 * self.pos.max()**2/2
        emin = self.v.min()**2/2 + self.cfg.omega**2 * self.pos.min()**2/2

        pmax = math.sqrt(2*emax)
        pmin = math.sqrt(2*emin)
        xmax = math.sqrt(2*emax / self.cfg.omega**2)
        xmin = math.sqrt(2*emin / self.cfg.omega**2)

        # Plot boundaries of energy planes
        phase = np.arange(0, 2*math.pi, 0.05)
        outerp, outerx = pmax*np.sin(phase), xmax*np.cos(phase)
        innerp, innerx = pmin*np.sin(phase), xmin*np.cos(phase)

        ax.plot(outerx, outerp, color='gray')
        ax.plot(innerx, innerp, color='gray')

        writer = None
        if dst:
            writer = matplotlib.animation.FFMpegWriter(fps=fps)
            writer.setup(fig, dst)

        def init():
            return scat,
        def update(newdata):
            self.increment(dt)
            scat.set_offsets(np.array([self.pos, self.v]).T)
            if writer:
                writer.grab_frame()
            return scat,

        times = np.arange(0, n*dt, dt)
        anim = FuncAnimation(fig, update, frames=times,
                             init_func=init, blit=False, interval=1/float(fps) * 1000,
                             repeat=False)
        plt.show()
        if writer:
            writer.finish()

sim = Simulator(n=500, deltav=1)
sim.animate(dt=0.01, fps=40, n=5000, dst='animation_b.mp4')