Simulate a rubber band (discretize and treat each segment as a damped spring)

as_springs.py

#!/usr/bin/env python

"""Simulate a rubber band.

Does so by approximating it as a series of connected springs.
"""

from __future__ import print_function

import argparse

from matplotlib import animation
import matplotlib.pyplot as plt
import numpy as np
import seaborn
import six


_SPRING_CONSTANT_DEFAULT = 1.0
_DAMPING_COEFF_DEFAULT = 0.0
_TIMESTEPS_DEFAULT = 225
_INTERVAL_DEFAULT = 25


class RubberBand(object):
    """Simulates a circular rubber band.

    Does so via discretizing the rubber band into evenly sized springs
    and connecting them with nodes with "mass".

    Each point is connected to two springs, each of exerts force in the
    direction of the (straight) spring. In addition, there is a damping
    force proportional to a given nodes velocity.

    See: en.wikipedia.org/wiki/Harmonic_oscillator#Damped_harmonic_oscillator

    Args:
        num_pts (int): The number of points on the rubber band in the
            discretization (will also be the number of springs).
        stretched_radius (float): The radius of the rubber band when stretched
            at time zero.
        radius (Optional[float]): The radius of the rubber band at rest.
        mass (Optional[float]): The mass of each "point" at the ends of a
            spring.
        spring_constant (Optional[float]): The constant of proportionality for
            Hooke's law.
        damping_coeff (Optional[float]): The "viscous damping coefficient" that
            provides damping of the motion (towards equilibrium).
    """

    points = None
    velocities = None
    spring_length = None

    def __init__(self, num_pts, stretched_radius, radius=1.0,
                 mass=1.0, spring_constant=_SPRING_CONSTANT_DEFAULT,
                 damping_coeff=_DAMPING_COEFF_DEFAULT):
        self.num_pts = num_pts
        self.stretched_radius = stretched_radius
        self.radius = radius
        self.mass = mass
        self.spring_constant = spring_constant
        self.damping_coeff = damping_coeff
        # Set the computed state.
        self._set_initial_state()

    def _set_initial_state(self):
        """Set the points and velocities for the initial state.

        Also computes the length of each spring "at rest".

        Only intended to be used by the constructor.
        """
        # Initially there is no velocity.
        self.velocities = np.zeros((self.num_pts, 2), order='F')
        # Equally space the points around the circle.
        theta = np.linspace(0.0, 2 * np.pi, self.num_pts, endpoint=False)
        self.points = np.empty((self.num_pts, 2), order='F')
        self.points[:, 0] = self.stretched_radius * np.cos(theta)
        self.points[:, 1] = self.stretched_radius * np.sin(theta)
        # Compute the length of each spring "at rest".
        self.spring_length = (2 * np.pi * self.radius) / self.num_pts

    def _compute_hooke_forces(self):
        r"""Compute the force on each point due to Hooke's law.

        Consider three points :math:`p_0, p_1, p_2` that determine two
        springs. We want to compute the force on :math:`p_1` due to
        Hooke's law. We have two direction vectors **away** from
        :math:`p_1`: :math:`d_0 = p_0 - p_1` and :math:`d_1 = p_2 - p_1`
        which have lengths :math:`L_0, L_1`. Thus Hooke's law says the
        force will be in the directions with magnitude equal to the
        displacement from the ideal length :math:`\ell`:

        .. math::

           F_0 = k \frac{d_0}{L_0} (L_0 - \ell)
           F_1 = k \frac{d_1}{L_1} (L_1 - \ell)

        Returns:
            numpy.ndarray: The Nx2 array of forces.
        """
        forces = np.empty((self.num_pts, 2), order='F')
        for index in six.moves.xrange(self.num_pts):
            prev_i = (index - 1) % self.num_pts
            next_i = (index + 1) % self.num_pts
            prev_pt = self.points[prev_i, :]
            point = self.points[index, :]
            next_pt = self.points[next_i, :]

            direction0 = prev_pt - point
            length0 = np.linalg.norm(direction0, ord=2)
            if length0 == 0.0:
                raise ValueError('Zero length')
            # Normalize the direction vector.
            direction0 /= length0

            direction1 = next_pt - point
            length1 = np.linalg.norm(direction1, ord=2)
            if length1 == 0.0:
                raise ValueError('Zero length')
            # Normalize the direction vector.
            direction1 /= length1

            force0_mag = self.spring_constant * (length0 - self.spring_length)
            force0 = force0_mag * direction0
            force1_mag = self.spring_constant * (length1 - self.spring_length)
            force1 = force1_mag * direction1
            forces[index, :] = force0 + force1

        return forces

    def update_state(self, delta_t):
        """Update the positions and velocities via Forward Euler.

        Args:
            delta_t (float): The timestep used to update the position
                and velocities (we have a 2nd order system since Newton's
                2nd law gives us acceleration from forces).
        """
        hooke_force = self._compute_hooke_forces()
        damping_force = -self.damping_coeff * self.velocities
        acceleration = (hooke_force + damping_force) / self.mass

        # Update points base on CURRENT velocities.
        self.points += delta_t * self.velocities
        # Update velocities based on acceleration.
        self.velocities += delta_t * acceleration


class RubberBandAnimation(object):
    """Animates a simulated circular rubber band.

    Args:
        rubber_band (RubberBand): The rubber band being animated.
        delta_t (float): The timestep used to update the position
            and velocities (we have a 2nd order system since Newton's
            2nd law gives us acceleration from forces).
    """

    def __init__(self, rubber_band, delta_t):
        self.rubber_band = rubber_band
        self.delta_t = delta_t
        # Set up the plot-specific values.
        self.figure = plt.figure()
        self.axes = self.figure.gca()
        self.frame_num = -1  # Precedes frame 0.
        self.animated_line = None

    def initialize_plot(self):
        """Create the first frame for :meth:`animate`.

        Returns:
            matplotlib.lines.Line2D: The line that was created with all
            the points on the rubber band.
        """
        stacked = np.vstack([
            self.rubber_band.points, self.rubber_band.points[[0], :]])
        self.animated_line, = self.axes.plot(
            stacked[:, 0], stacked[:, 1], marker='o')

        # Set up the axes based on the ``stretched_radius``.
        max_coord = 1.125 * self.rubber_band.stretched_radius
        self.axes.axis('scaled')
        self.axes.set_xlim(-max_coord, max_coord)
        self.axes.set_ylim(-max_coord, max_coord)

        return self.animated_line,

    def update_line(self, frame_num):
        """Update a matplotlib line with the current points.

        Used by :meth:`animate`.

        First updates the points / velocities with one timestep of size
        ``delta_t``.

        Does so by adding the first point to the end, so that the points
        form a circle.

        Args:
            frame_num (int): The current frame number.

        Returns:
            matplotlib.lines.Line2D: The line that was updated with the
            current points on the rubber band.
        """
        if frame_num == self.frame_num + 1:
            self.frame_num = frame_num
        else:
            raise ValueError('Unexpected frame number.')

        self.rubber_band.update_state(self.delta_t)

        stacked = np.vstack([
            self.rubber_band.points, self.rubber_band.points[[0], :]])
        self.animated_line.set_xdata(stacked[:, 0])
        self.animated_line.set_ydata(stacked[:, 1])

        return self.animated_line,

    def animate(self, steps=_TIMESTEPS_DEFAULT,
                interval=_INTERVAL_DEFAULT, filename=None):
        """Create an animation for the motion of the current rubber band.

        Args:
            steps (Optional[int]): The number of time steps to simulate.
            interval (Optional[int]): Delay between frames in milliseconds.
            filename (Optional[str]): Name of MP4 file to save plot in. If
                not specified, will just display the animation.
        """
        line_animation = animation.FuncAnimation(
            self.figure, self.update_line, frames=steps,
            init_func=self.initialize_plot,
            interval=interval, blit=True, repeat=False)

        # Plot or save.
        if filename is None:
            plt.show()
        else:
            print('Saving {}...'.format(filename))
            fps = int(1000.0 / interval)
            line_animation.save(
                filename=filename, fps=fps, dpi=200, writer='ffmpeg')
            print('Saved {}'.format(filename))


def get_parsed_args():
    """Get arguments parsed from the command line.

    Returns:
        argparse.Namespace: Parsed arguments.
    """
    parser = argparse.ArgumentParser(
        description='Simulate a rubber band.',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--num-points', dest='num_pts',
                        type=int, default=16)
    parser.add_argument('--stretched-radius', dest='stretched_radius',
                        type=float, default=1.5)
    parser.add_argument('--spring-constant', dest='spring_constant',
                        type=float, default=_SPRING_CONSTANT_DEFAULT)
    parser.add_argument('--damping-coeff', dest='damping_coeff',
                        type=float, default=_DAMPING_COEFF_DEFAULT)
    parser.add_argument('--delta-t', dest='delta_t',
                        type=float, default=0.125)
    parser.add_argument('--steps', type=int, default=_TIMESTEPS_DEFAULT)
    parser.add_argument('--interval', type=int, default=_INTERVAL_DEFAULT)
    parser.add_argument('--filename')
    return parser.parse_args()


def main():
    """Run the simulation based on command line arguments."""
    args = get_parsed_args()
    rubber_band = RubberBand(
        args.num_pts, args.stretched_radius,
        spring_constant=args.spring_constant,
        damping_coeff=args.damping_coeff)
    rubber_band_anim = RubberBandAnimation(rubber_band, args.delta_t)
    rubber_band_anim.animate(
        steps=args.steps, interval=args.interval, filename=args.filename)


if __name__ == '__main__':
    main()

commands.txt

./as_springs.py --filename rubber_band_wo_damping.mp4
./as_springs.py --steps 500 --filename rubber_band_wo_damping_long.mp4
./as_springs.py --damping-coeff 0.5 --filename rubber_band_w_damping.mp4

rubber_band_w_damping.mp4

rubber_band_wo_damping.mp4

rubber_band_wo_damping_long.mp4