## File: barbell.R
## This file contains a simulation of a falling barbell. Dropping a barbell (or a needle)
## confined to 2 dimensions is similar to a coin toss. A barbell has 2 ends, marked A and B.
## There are two equilibria: one with A to the right of B and one with A to the left of B.
## Thus, there are two outcomes. The barbell loses energy due to air resistance. The
## probability of one of the 2 equilibrium states depends on the initial condition.
## If the altitude is too small, and initial angular velocity is small, there is not
## enough energy to "flip" and the barbell ends up in the same state as it started.
##
## On a technical level, this implementation uses "S4 classes", one of several ways
## to use object oriented programming in R.
##
library(animation)

## State of the barbell
barbell <- setClass("barbell",
                    slots = c(
                      m="numeric",      # mass
                      I="numeric",      # moment of inertia
                      l="numeric",      # length
                      x="numeric",      # distance from floor
                      p="numeric",      # momentum
                      theta="numeric",  # angle with horizonal direction
                      omega="numeric",  # angular momentum
                      g="numeric",      # gravity constant
                      A="numeric",      # left end
                      B="numeric",      # right end
                      HookeConstant="numeric", # floor Hooke constant
                      AirResistance="numeric"  # air friction factor
                      ),
                    prototype = list(
                      m=1,
                      I=numeric(0),
                      l=1,
                      x=5,
                      p=0,
                      theta=0,
                      omega=0,
                      g=1,
                      HookeConstant=1,
                      AirResistance=0
                      )
                    );

setGeneric("update",def=function(.Object){stop("Not implemented")});
setMethod("update",
          signature=c("barbell"),
          definition=function(.Object) {
            r <- c(cos(.Object@theta),sin(.Object@theta));
            .Object@A = c(0,.Object@x)+.Object@l*r;
            .Object@B = c(0,.Object@x)-.Object@l*r;
            .Object
          });


## Initializes a barbell, calls update
setMethod("initialize", 
          signature = c("barbell"),
          definition = function(.Object,...) {
            .Object <- callNextMethod();
            .Object@I <- 2*.Object@m*.Object@l^2;
            .Object <- update(.Object);
            .Object
          });


## We use Lagrangian mechanics.
## We use a potential to represent the forces.
## L=T-V
## H=T+V is Hamiltonian
## T=1/2*p^2/m + 1/2*omega^2/I
setGeneric("draw",def=function(.Object){stop("Not implemented")})
setMethod("draw",
          signature=c("barbell"),
          definition=function(.Object){
            A <- .Object@A;
            B <- .Object@B;
            lines(c(A[1],B[1]),c(A[2],B[2]),col="orange");
            points(c(A[1],B[1]),c(A[2],B[2]),col=c("red","blue"),cex=5);
            text(A[1],A[2],"A",adj=0.5);
            text(B[1],B[2],"B",adj=0.5);
            title("Bouncing Barbell");
            .Object
          });

## Implements evolution according to the equations of motion.
## The derivation of the equations is in barbell.tex/barbell.pdf.
## The integration scheme is the forward Euler scheme.
setGeneric("evolve",def=function(.Object,...){stop("Not implemented")})
setMethod("evolve",
          signature=c("barbell"),
          definition=function(.Object,timeStep=.1){
            hA <- .Object@x+.Object@l*sin(.Object@theta);
            hB <- .Object@x-.Object@l*sin(.Object@theta);
            fA <- ifelse(hA>0,-.Object@m*.Object@g,-.Object@HookeConstant*hA);
            fB <- ifelse(hB>0,-.Object@m*.Object@g,-.Object@HookeConstant*hB);
            .Object@p <- .Object@p + timeStep*(fA+fB)-
              .Object@p*.Object@AirResistance/2/.Object@m;
            .Object@omega <- .Object@omega+timeStep*(fA-fB)*.Object@l*cos(.Object@theta)-
              .Object@omega*.Object@AirResistance/.Object@I;
            .Object@x <- .Object@x+timeStep*.Object@p/(2*.Object@m);
            .Object@theta <- .Object@theta+timeStep*.Object@omega/.Object@I;
            .Object <- update(.Object);
            .Object
          });


## Set delay between frames when replaying
ani.options(interval=.02);

## Set up a vector of colors for use below
col.range <- heat.colors(15);

## A bouncing barbell simulation
animate.barbell <- function() {
  saveHTML({
    timeStep <- .2;
    xlim <- c(-10,10);
    ylim <- c(-5,15);
    my.barbell <- barbell(omega=10,
                          p=0,
                          x=12,
                          l=3,
                          HookeConstant=60,
                          AirResistance=.08);
    for(i in 1:500) {
      plot.new();
      plot.window(xlim,ylim);
      abline(-1,0,col="red");
      draw(my.barbell);
      my.barbell <- evolve(my.barbell,timeStep);
    }
  })
}


test.draw <- function() {
  my.barbell <- barbell(theta=pi/6);
  xlim <- c(-10,10);
  ylim <- c(-5,15);
  plot.new();
  plot.window(xlim,ylim);
  abline(0,0,col="red");
  draw(my.barbell);
}

## test.draw()