Animated oscillations with x(t) plot. phys.pro/problems/1328

1328.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.lines as lines
from matplotlib.patches import FancyArrowPatch

def generate_vertical_spring(start, end, loops=10, points_per_loop=100):
    '''
    Generate coordinates for a vertical spring from start to end with a specified number of loops.
    
    Parameters:
    start (tuple): Coordinates (x, y) of the start point.
    end (tuple): Coordinates (x, y) of the end point.
    loops (int): Number of loops in the spring.
    points_per_loop (int): Number of points per loop.
    
    Returns:
    list: List of (x, y) coordinates representing the spring.
    '''
    x_start, y_start = start
    x_end, y_end = end
    
    total_length = y_end - y_start
    
    t = np.linspace(0, 2 * np.pi * loops, loops * points_per_loop)
    x = 0.07 * np.sin(t)  # Adjust amplitude (0.1) to control spring width
    y = np.linspace(y_start, y_end, loops * points_per_loop)
    
    # Shift x to start at the given x_start
    x += x_start
    
    path = list(zip(x, y))
    return path

def plot_spring(path, start, end, amplitude, w0, ti=None):
    '''
    Plot the spring given a list of (x, y) coordinates and plot the oscillation.
    
    Parameters:
    path (list): List of (x, y) coordinates representing the spring.
    start (tuple): Coordinates (x, y) of the start point.
    end (tuple): Coordinates (x, y) of the end point.
    amplitude (float): Amplitude of oscillations in meters.
    w0 (float): Angular frequency of oscillations.
    '''
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6), gridspec_kw={'width_ratios': [0.5, 1]}, sharey=True)
    fig.subplots_adjust(wspace=-0.54)
    
    # Plot the spring
    ax1.plot(*zip(*path), lw=2)
    
    # Add ceiling
    ax1.plot([-0.3, 0.3], [-0.007, -0.007], 'k', lw=4, clip_on=False)

    # Add weight (circle) at the bottom
    ax1.plot(end[0], end[1], '.', color='#AA6666', ms=50, zorder=50)
    
    # Set limits for spring plot
    ax1.set_ylim([1.5, 0])
    ax1.set_xlim([-0.4, 0.4])

    # Add vertical axis with arrow
    arrowstyle_y = "-|>,head_length=0.9,head_width=0.3"
    arrow_y = FancyArrowPatch( (-0.3, 0), (-0.3, 1.3),
                              clip_on=False,
                              arrowstyle=arrowstyle_y,
                              mutation_scale=10,
                              shrinkA=0,
                              shrinkB=0,
                              lw=2, color='black', zorder=5)
    ax1.add_patch(arrow_y)
    ax1.text(-0.3, 1.3, '$x$', verticalalignment='top', horizontalalignment='center')

    dx = 0.312
    dy = l0 + xp
    
    # Plot horizontal l0 line
    ax1.plot([-dx, 0.3], [l0, l0], 'k--', lw=1.3)

    xarr=0.201
    # Plot l0 arrow
    arrowstyle_y = "<|-|>,head_length=0.7,head_width=0.2"
    arrow_y = FancyArrowPatch( (xarr, 0), (xarr, l0),
                              clip_on=False,
                              arrowstyle=arrowstyle_y,
                              mutation_scale=10,
                              shrinkA=0,
                              shrinkB=0,
                              lw=1.5, color='black', zorder=5)
    ax1.add_patch(arrow_y)
    ax1.text(0.212, l0/2, r'$\ell_0$', verticalalignment='center')

    # Plot xp arrow
    arrow_y = FancyArrowPatch( (xarr, l0), (xarr, l0+xp),
                              clip_on=False,
                              arrowstyle=arrowstyle_y,
                              mutation_scale=10,
                              shrinkA=0,
                              shrinkB=0,
                              lw=1.5, color='black', zorder=5)
    ax1.add_patch(arrow_y)
    ax1.text(0.212, l0+xp/1.8, r'$x_p$', verticalalignment='center')

    # Plot A arrow
    arrow_y = FancyArrowPatch( (xarr, l0+xp), (xarr, l0+xp+amplitude),
                              clip_on=False,
                              arrowstyle=arrowstyle_y,
                              mutation_scale=10,
                              shrinkA=0,
                              shrinkB=0,
                              lw=1.5, color='black', zorder=5)
    ax1.add_patch(arrow_y)
    ax1.text(0.212, l0+xp+amplitude/2, r'$A$', verticalalignment='center')


    shift = 0
    T = 2 * np.pi / w0  # Period of oscillation

    # Plot horizontal amplitudes and equilibrium always
    ax1.plot([-dx, 0.3], [dy-amplitude, dy-amplitude], 'k--', lw=1.3)
    ax1.plot([-dx, 0.3], [dy+amplitude, dy+amplitude], 'k--', lw=1.3)
    ax1.plot([-dx, 0.3], [dy, dy], 'k--', lw=1.3)

    # Fake x_p-A for maintain plot size
    ax1.text(-dx*1.01, l0, '$x_p-A$', verticalalignment='center', horizontalalignment='right', color='white', zorder=-99)


    if ti < 2*T:
        # Plot horizontal 0 line
        ax1.plot([-dx, 0.3], [l0+xp, l0+xp], 'k--', lw=1.3)

        # Add text for horizontal amplitudes
        ax1.text(-dx*1.01, l0+xp, '$0$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp-amplitude, '$-A$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp+amplitude, '$A$', verticalalignment='center', horizontalalignment='right')

    elif ti >= 2*T and ti < 3*T:
        shift = - amplitude * (ti - 2 * T) / T

        # Plot horizontal 0 line
        ax1.plot([-dx, 0.3], [dy+shift, dy+shift], 'k--', lw=1.3)
        ax1.text(-dx*1.01, dy+shift, '$0$', verticalalignment='center', horizontalalignment='right')

    elif ti >= 3*T and ti < 4*T:

        # Add text for horizontal amplitudes
        ax1.text(-dx*1.01, l0+xp-amplitude, '$0$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp, '$A$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp+amplitude, '$2A$', verticalalignment='center', horizontalalignment='right')

    elif ti >= 4*T and ti < 5*T:
        shift =  - (xp-amplitude) * (ti - 4 * T) / T

        # Plot horizontal 0 line
        ax1.plot([-dx, 0.3], [dy-amplitude+shift, dy-amplitude+shift], 'k--', lw=1.3)
        ax1.text(-dx*1.01, dy-amplitude+shift, '$0$', verticalalignment='center', horizontalalignment='right')

    elif ti >= 5*T and ti < 6*T:
        # Add text for horizontal amplitudes
        ax1.text(-dx*1.01, l0, '$0$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp-amplitude, '$x_p-A$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp, '$x_p$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp+amplitude, '$x_p+A$', verticalalignment='center', horizontalalignment='right')

    elif ti >= 6*T and ti < 7*T:
        shift = xp * (ti - 6 * T) / T
        # Plot horizontal 0 line
        ax1.plot([-dx, 0.3], [l0+shift, l0+shift], 'k--', lw=1.3)
        ax1.text(-dx*1.01, l0+shift, '$0$', verticalalignment='center', horizontalalignment='right')

    elif ti == 7*T:
        # Add text for horizontal amplitudes
        ax1.text(-dx*1.01, l0+xp, '$0$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp-amplitude, '$-A$', verticalalignment='center', horizontalalignment='right')
        ax1.text(-dx*1.01, l0+xp+amplitude, '$A$', verticalalignment='center', horizontalalignment='right')

    ax1.axis('off')
        


    #######################
    ##### RIGHT PLOT ######
    #######################
        
    # Plot the x = A * cos(w0 * t) function on the second axis
    T = 2 * np.pi / w0  # Period of oscillation
    t = np.linspace(0, T, 500)
    x = - amplitude * np.cos(w0 * t) + dy
    y = t  # Use t for horizontal progression
    
    ax2.plot(y, x, lw=2)
    
    # Make original axes invisible
    ax2.spines['top'].set_color('none')
    ax2.spines['right'].set_color('none')
    ax2.spines['bottom'].set_color('none')
    ax2.spines['left'].set_color('none')
    
    ax2.set_yticks([])

    # Add horizontal axis with arrow
    dx=-0.02
    # Fake x_p-A for maintaine fig size
    ax2.text(dx, l0, '$x_p-A$', horizontalalignment='right', verticalalignment='center', color='white', zorder=-99)

    arrowstyle_x = "-|>,head_length=0.9,head_width=0.3"

    if ti < 2*T:
        arrow_x = FancyArrowPatch( (0, dy), (T*1.1, dy),
                                  clip_on=False,
                                  arrowstyle=arrowstyle_x,
                                  mutation_scale=10,
                                  shrinkA=0,
                                  shrinkB=0,
                                  lw=2, color='black', zorder=5)
        ax2.add_patch(arrow_x)
        ax2.text(T*1.11, dy, '$t$', horizontalalignment='center', verticalalignment='bottom')

        ax2.text(dx, dy-amplitude, '$-A$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, dy, '$0$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, dy+amplitude, '$A$', horizontalalignment='right', verticalalignment='center')

        # x(t) solution
        ax2.text(0.05, 0.131, r'{', fontsize=40)
        ax2.text(0.1, 0.07, r'$m\ddot{x}\;=\; mg - k(x+x_p)$')
        ax2.text(0.1, 0.15, r'$mg\;=\; kx_p$')
        #ax2.text(0.1, 0.25, r'$\ddot{x} + \dfrac{k}{m}x=0$')
        ax2.text(0.1, 0.25, r'$\ddot{x} + (k/m)\cdot x=0$')
        ax2.text(0.1, 0.35, r'$x\,(t)\;=\; -A\cos\omega_0t$')


    elif ti >= 2*T and ti < 3*T:
        arrow_x = FancyArrowPatch((0, dy+shift), (T*1.1, dy+shift),
                                  clip_on=False,
                                  arrowstyle=arrowstyle_x,
                                  mutation_scale=10,
                                  shrinkA=0,
                                  shrinkB=0,
                                  lw=2, color='black', zorder=5)
        ax2.add_patch(arrow_x)
        ax2.text(T*1.11, dy+shift, '$t$', horizontalalignment='center', verticalalignment='bottom')
        ax2.text(dx, dy+shift, '$0$', horizontalalignment='right', verticalalignment='center')

        # x(t) solution
        ax2.text(0.05, 0.131, r'{', fontsize=40)
        ax2.text(0.1, 0.07, r'$m\ddot{x}\;=\; mg - k(x+x_p-\Delta x)$')
        ax2.text(0.1, 0.15, r'$mg\;=\; kx_p$')
        ax2.text(0.1, 0.25, r'$\ddot{x} + (k/m)\cdot x - k\Delta x/m \; = \; 0$')
        ax2.text(0.1, 0.35, r'$x\,(t)\;=\; -A\cos\omega_0t + \Delta x$')

    elif ti >= 3*T and ti < 4*T:
        arrow_x = FancyArrowPatch((0, dy-amplitude), (T*1.1, dy-amplitude),
                                  clip_on=False,
                                  arrowstyle=arrowstyle_x,
                                  mutation_scale=10,
                                  shrinkA=0,
                                  shrinkB=0,
                                  lw=2, color='black', zorder=5)
        ax2.add_patch(arrow_x)
        ax2.text(T*1.11, dy-amplitude, '$t$', horizontalalignment='center', verticalalignment='bottom')

        ax2.text(dx, dy-amplitude, '$0$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, dy, '$A$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, dy+amplitude, '$2A$', horizontalalignment='right', verticalalignment='center')

        # x(t) solution
        ax2.text(0.05, 0.131, r'{', fontsize=40)
        ax2.text(0.1, 0.07, r'$m\ddot{x}\;=\; mg - k(x+x_p-A)$')
        ax2.text(0.1, 0.15, r'$mg\;=\; kx_p$')
        ax2.text(0.1, 0.25, r'$\ddot{x} + (k/m)\cdot x - kA/m \; = \; 0$')
        ax2.text(0.1, 0.35, r'$x\,(t)\;=\; -A\cos\omega_0t + A$')

    elif ti >= 4*T and ti < 5*T:
        shift =  - (xp-amplitude) * (ti - 4 * T) / T
        # Move horizontal axis to l0
        arrow_x = FancyArrowPatch((0, dy-amplitude+shift), (T*1.1, dy-amplitude+shift),
                                  clip_on=False,
                                  arrowstyle=arrowstyle_x,
                                  mutation_scale=10,
                                  shrinkA=0,
                                  shrinkB=0,
                                  lw=2, color='black', zorder=5)
        ax2.add_patch(arrow_x)
        ax2.text(T*1.11, dy-amplitude+shift, '$t$', horizontalalignment='center', verticalalignment='bottom')
        ax2.text(dx, dy-amplitude+shift, '$0$', horizontalalignment='right', verticalalignment='center')

        # x(t) solution
        ax2.text(0.05, 0.131, r'{', fontsize=40)
        ax2.text(0.1, 0.07, r'$m\ddot{x} \; = \; mg - k(x+x_p - \Delta x)$')
        ax2.text(0.1, 0.15, r'$mg \; = \; kx_p$')
        ax2.text(0.1, 0.25, r'$\ddot{x} + (k/m)\cdot x - k\Delta x/m \; = \; 0$')
        ax2.text(0.1, 0.35, r'$x\,(t)\;=\; -A\cos\omega_0t + \Delta x$')

    elif ti >= 5*T and ti < 6*T:
        arrow_x = FancyArrowPatch((0, l0), (T*1.1, l0),
                                  clip_on=False,
                                  arrowstyle=arrowstyle_x,
                                  mutation_scale=10,
                                  shrinkA=0,
                                  shrinkB=0,
                                  lw=2, color='black', zorder=5)
        ax2.add_patch(arrow_x)
        ax2.text(T*1.11, l0, '$t$', horizontalalignment='center', verticalalignment='bottom')
        ax2.text(dx, l0, '$0$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, l0+xp-amplitude, '$x_p-A$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, l0+xp, '$x_p$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, l0+xp+amplitude, '$x_p+A$', horizontalalignment='right', verticalalignment='center')

        # x(t) solution
        ax2.text(0.05, 0.131, r'{', fontsize=40)
        ax2.text(0.1, 0.07, r'$m\ddot{x}\;=\; mg - kx$')
        ax2.text(0.1, 0.15, r'$mg\;=\; kx_p$')
        ax2.text(0.1, 0.25, r'$\ddot{x} + (k/m)\cdot x - g \; = \; 0$')
        ax2.text(0.1, 0.35, r'$x\,(t)\;=\; -A\cos\omega_0t + x_p$')

    elif ti >= 6*T and ti < 7*T:
        shift =  xp * (ti - 6 * T) / T
        arrow_x = FancyArrowPatch((0, l0+shift), (T*1.1, l0+shift),
                                  clip_on=False,
                                  arrowstyle=arrowstyle_x,
                                  mutation_scale=10,
                                  shrinkA=0,
                                  shrinkB=0,
                                  lw=2, color='black', zorder=5)
        ax2.add_patch(arrow_x)
        ax2.text(T*1.11, l0+shift, '$t$', horizontalalignment='center', verticalalignment='bottom')
        ax2.text(dx, l0+shift, '$0$', horizontalalignment='right', verticalalignment='center')

        # x(t) solution
        ax2.text(0.05, 0.131, r'{', fontsize=40)
        ax2.text(0.1, 0.07, r'$m\ddot{x}\;=\; mg - k(x+x_p-\Delta x)$')
        ax2.text(0.1, 0.15, r'$mg\;=\; kx_p$')
        ax2.text(0.1, 0.25, r'$\ddot{x} + (k/m)\cdot x=0$')
        ax2.text(0.1, 0.35, r'$x\,(t)\;=\; -A\cos\omega_0t + \Delta x$')

    elif ti == 7*T:
        # Add text for A
        ax2.text(dx, dy-amplitude, '$-A$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, dy, '$0$', horizontalalignment='right', verticalalignment='center')
        ax2.text(dx, dy+amplitude, '$A$', horizontalalignment='right', verticalalignment='center')

        # Arrow
        arrow_x = FancyArrowPatch( (0, dy), (T*1.1, dy),
                                  clip_on=False,
                                  arrowstyle=arrowstyle_x,
                                  mutation_scale=10,
                                  shrinkA=0,
                                  shrinkB=0,
                                  lw=2, color='black', zorder=5)
        ax2.add_patch(arrow_x)
        ax2.text(T*1.11, dy, '$t$', horizontalalignment='center', verticalalignment='bottom')

        # x(t) solution
        ax2.text(0.05, 0.131, r'{', fontsize=40)
        ax2.text(0.1, 0.07, r'$m\ddot{x}\;=\; mg - k(x+x_p)$')
        ax2.text(0.1, 0.15, r'$mg\;=\; kx_p$')
        ax2.text(0.1, 0.25, r'$\ddot{x} + (k/m)\cdot x=0$')
        ax2.text(0.1, 0.35, r'$x\,(t)\;=\; -A\cos\omega_0t$')


    ax2.set_aspect(1)

    # Add vertical axis with arrow
    arrowstyle_y = "-|>,head_length=0.9,head_width=0.3"
    arrow_y = FancyArrowPatch( (0, 0), (0, 1.3),
                              clip_on=False,
                              arrowstyle=arrowstyle_y,
                              mutation_scale=10,
                              shrinkA=0,
                              shrinkB=0,
                              lw=2, color='black', zorder=5)
    ax2.add_patch(arrow_y)
    ax2.text(0, 1.3, '$x$', horizontalalignment='center', verticalalignment='top')

    # Add vertical dashed lines at T/4, T/2, 3T/4, T
    lw=1.3

    ax2.plot( [T/4, T/4], [dy-amplitude, dy+amplitude+0.02], color='black', linestyle='--', lw=lw)
    ax2.plot( [T/2, T/2], [dy-amplitude, dy+amplitude+0.02], color='black', linestyle='--', lw=lw)
    ax2.plot( [3*T/4, 3*T/4], [dy-amplitude, dy+amplitude+0.02], color='black', linestyle='--', lw=lw)
    ax2.plot( [T, T], [dy-amplitude, dy+amplitude+0.02], color='black', linestyle='--', lw=lw)

    # Add horizontal dashed lines for -A, A
    ax2.plot([0,T], [dy+amplitude, dy+amplitude], color='black', linestyle='--', lw=1.3, clip_on=False)
    ax2.plot([0,T], [dy, dy], color='black', linestyle='--', lw=1.3, clip_on=False)
    ax2.plot([0,T], [dy-amplitude, dy-amplitude], color='black', linestyle='--', lw=1.3, clip_on=False)

    # Indicate phase with red dot
    ti_mod = np.mod(ti, T)
    ax2.plot(ti_mod, dy - amplitude * np.cos(w0 * ti_mod), '.', color='#AA6666', ms=15)
    
    # Set custom ticks without original axes lines
    #ax2.set_xticks([0, T / 4, T / 2, 3 * T / 4, T], [0, '$T/4$','$T/2$','$3T/4$', '$T$'])
    ax2.set_xticks([T / 4, T / 2, 3 * T / 4, T], ['$T/4$','$T/2$','$3T/4$', '$T$'])
    ax2.tick_params(axis='x', length=0, width=0)
    ax2.tick_params(axis='both', which='major', pad=-82)
    
    # Set limits for the oscillation plot
    ax2.set_xlim([0, T*1.05])
    #ax2.set_ylim([0,2])
    #ax2.xaxis.tick_top()
    
    
    plt.tight_layout()
    #plt.savefig('1328/{:.3f}.png'.format(round(ti,3)), dpi=286, bbox_inches='tight')
    plt.savefig('1328/{:.3f}.png'.format(round(ti,3)), dpi=289, bbox_inches='tight')
    plt.close()
    #plt.show()

# Define parameters for the spring and oscillation
k = 12.5  # N/m, spring constant
m = 0.5  # kg, mass of the weight
g = 10   # m/s^2, gravitational acceleration
w0 = np.sqrt(k / m)  # Angular frequency
amplitude = 0.25  # Amplitude of oscillation in meters (20 cm)
l0 = 0.5
T = 2*np.pi / w0
times = np.linspace(0,7*T, 1200)


plt.rcParams.update({
    "text.usetex": False,
    "font.size": 20,
    "mathtext.fontset": "custom",
    "mathtext.it": "STIX Two Text:italic",
    "mathtext.rm": "STIX Two Text",
    "mathtext.sf": "STIX Two Text",
    "font.sans-serif": "STIX Two Text",
})


#ti = T/4
for ti in times:
    x = amplitude * np.cos(w0*ti)
    xp = m * g / k  # length of the spring in equilibrium position

    start_point = (0, 0)
    end_point = (0, l0 - x + xp)
    # Generate the spring path
    spring_path = generate_vertical_spring(start_point, end_point, loops=4)
    plot_spring(spring_path, start_point, end_point, amplitude, w0, ti = ti)