Example: Bouncing ball Euler and RK4

gistfile1.lua

--------------------------------------------------------------------------------
-- License
-- =============================================================================
-- Copyright (c) 2015 Wouter Lindenhof <wouterlindenhof@gmail.com>
--
-- Permission is hereby granted, free of charge, to any person obtaining a copy 
-- of this software and associated documentation files (the "Software"), to deal
-- in the Software without restriction, including without limitation the rights 
-- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 
-- copies of the Software, and to permit persons to whom the Software is 
-- furnished to do so, subject to the following conditions:
-- 
-- The above copyright notice and this permission notice shall be included in 
-- all copies or substantial portions of the Software.

-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
-- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
-- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
-- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
-- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-- SOFTWARE.
--
-- Introduction
-- =============================================================================
-- 
-- This file is a complete example of why euler is bad for intergration 
-- calculation and why Runge-Kutta is better. It has been written so that it is
-- easy to read and understand. 
-- 
-- We do require that the reader is familiar with both basic vector math and 
-- programming. The code has been written in Lua and if you want to run it you
-- can download LÖVE from https://www.love2d.org and use that to run the 
-- example.
-- 
-- How to read this file
-- =============================================================================
-- 
-- The best way to read this file is from top to bottom. We start with the 
-- simple `scene_setup` which is followed by `update_ball`. At the end of the 
-- file you will find some framework code and utility classes.
-- 
--------------------------------------------------------------------------------


--------------------------------------------------------------------------------
-- These variables are of importance of the logic following, but shouldn't be 
-- changed since they depend on the size of the window and resolution. They are
-- only shown so that you the reader are aware that they exist.
--------------------------------------------------------------------------------
balls        = {}   -- Table in which we store the balls
show_intro   = true -- To start the simulation the user needs to press a key.
GRAVITY      = 300  -- How many pixels should the ball fall per second?
FLOOR_POS_Y  = 550  -- Where is the floor drawn?
BALL_RADIUS  = 16   -- Size of the ball
START_Y      = 120  -- At what height do we place the balls?


--------------------------------------------------------------------------------
-- Setup the scene up with 4 balls. Each ball starts the same and should bounce
-- the same where it not for 
--------------------------------------------------------------------------------
function setup_scene()
  -- We create four balls and only pass properties that are different for each 
  -- ball. The `create_ball` will create an object and fill all the values.
  -- 
  -- Properties passed to create_ball:
  -- 1) The x position on the screen (otherwise they overlap)
  -- 2) The interpolation (since we want to demonstrate rk4 versus euler)
  -- 3) The amount of simulation frames
  --
  local b1 = create_ball(160, 'linear',  1)
  local b2 = create_ball(320, 'linear', 10)
  local b3 = create_ball(480, 'rk4',     1)
  local b4 = create_ball(640, 'rk4',    10)

  table.insert(balls, b1)
  table.insert(balls, b2)
  table.insert(balls, b3)
  table.insert(balls, b4)
end


--------------------------------------------------------------------------------
-- Create a table that represents a ball and then return it. 
--------------------------------------------------------------------------------
function create_ball(pos_x, interpolation, sim_step)
  local ball = {}

  ball.pos           = Vec2.new(pos_x, START_Y) -- Position
  ball.vel           = Vec2.new()               -- Velocity
  ball.force         = Vec2.new(0, GRAVITY)     -- Acceleration (gravity)
  ball.interpolation = interpolation            -- Method of interpolation
  ball.sim_step      = sim_step                 -- Simulation steps per frame

  return ball
end


--------------------------------------------------------------------------------
-- update_ball is called once for every ball per draw frame. 
--------------------------------------------------------------------------------
function update_ball(ball, dt)
  -- To improve precision of the simulation we can simulate smaller steps. 
  -- Smaller steps improve the simulation.
  local sim_dt = (dt / ball.sim_step)
  for step=1, ball.sim_step do
    simulate_ball(ball, sim_dt)
  end
end


--------------------------------------------------------------------------------
-- Perform a single simulation step of the ball. Within a frame this function 
-- can be called multiple times to increase the precision. The `dt` is the 
-- amount of time that must be simulated.
--------------------------------------------------------------------------------
function simulate_ball(ball, dt)
  -- Retrieve the function we use to interpolate and then perform the 
  -- interpolation. The values returned from the function are directly stored
  -- in the ball.
  local f = get_interpolation(ball.interpolation)
  ball.pos, ball.vel = f(ball.pos, ball.vel, ball.force, dt)

  -- Check if we need to revert the velocity if the ball hits the floor. This 
  -- function does not correct the position of the ball (allowing the ball to 
  -- fall through the floor) as we are more concerned with preserving energy.
  bounce_check(ball)
end


--------------------------------------------------------------------------------
-- The bounce_check reverts the ball if it hits the floor and is falling. If the
-- ball is not falling then we do not revert the ball because we can only hit 
-- the floor when falling.
--------------------------------------------------------------------------------
function bounce_check(ball)
  local hits_floor = (ball.pos.y + BALL_RADIUS) > FLOOR_POS_Y;
  local is_falling = ball.vel.y > 0
  if hits_floor and is_falling then
    ball.vel.y = -ball.vel.y -- revert the velocity
  end
end


--------------------------------------------------------------------------------
-- Based on the name it will returns a function that can be used to calculate 
-- the future position and velocity of an object.
--------------------------------------------------------------------------------
function get_interpolation(name)
  if name == 'linear' then
    return euler_intergration
  elseif name == 'rk4' then
    return rk4_intergration
  end
end


--------------------------------------------------------------------------------
-- The basic and most error prone intergration function for position and 
-- velocity.
--------------------------------------------------------------------------------
function euler_intergration(pos, vel, acc, dt)
  vel = vel + (acc * dt)
  pos = pos + (vel * dt)

  return pos, vel
end


--------------------------------------------------------------------------------
-- Runge-Kutta fourth order is another intergration method which is far less 
-- error prone then euler.
-- What basically happens is that four samples are taken of both position and 
-- velocity. Once at the begin of the frame, once at the end of the frame and 
-- two at the midpoint of the frame (although the initial conditions will be 
-- different). From these samples an weighted average is calculated which is
-- used to increment the position and velocity.
-- 
-- Of course, it is a bit more complex then that. The best place to start is the
-- average function and ask yourself how the weight is divided and why. Once 
-- that is clear look at how each k-value depends on each other and then you can
-- attempt to tackle `rk4eval`, which although the most simple looking part is 
-- actually the most complex part. That was at least the way I finally learned 
-- it.
--------------------------------------------------------------------------------
function rk4_intergration(pos, vel, acc, dt)
  local rk4eval = function (pos, vel, acc, dt, ki)
    local sx = pos + ki.x * dt
    local sv = vel + ki.v * dt
    return { x = sv, v = acc}
  end

  local k0 = { x = Vec2.new(), v = Vec2.new() }
  local k1 = rk4eval(pos, vel, acc, dt * 0.0, k0) -- Begin of the frame
  local k2 = rk4eval(pos, vel, acc, dt * 0.5, k1) -- Midpoint 
  local k3 = rk4eval(pos, vel, acc, dt * 0.5, k2) -- Midpoint
  local k4 = rk4eval(pos, vel, acc, dt * 1.0, k3) -- End of the frame

  -- Average the results
  local dxdt = (1.0 / 6.0) * (k1.x + 2.0 * (k2.x + k3.x) + k4.x );
  local dvdt = (1.0 / 6.0) * (k1.v + 2.0 * (k2.v + k3.v) + k4.v );

  -- Integrate the averages in the final results.
  pos = pos + dxdt * dt
  vel = vel + dvdt * dt

  return pos, vel
end


--------------------------------------------------------------------------------
-- LÖVE
-- =============================================================================
--
-- The following code section is code that is primarily used for the drawing and
-- calling the actual update function of the balls.
-- 
--------------------------------------------------------------------------------


--------------------------------------------------------------------------------
-- Callback that is called when love is loaded. Sets the resolution and then 
-- calls the `setup_scene` which is used as entry point of this example.
--------------------------------------------------------------------------------
function love.load()
  love.window.setMode(800, 600, {resizable=false, vsync=false, highdpi=true})
  setup_scene()
end


--------------------------------------------------------------------------------
-- Callback that is called every frame. It calls the `update_ball` which is 
-- considerd the second entry point of this example.
--------------------------------------------------------------------------------
function love.update( dt )
  -- If the `show_intro` is active then don't update anything.
  if show_intro then 
    return
  end

  for i,v in pairs(balls) do
    update_ball(v, dt)
  end
end


--------------------------------------------------------------------------------
-- Callback that is called when a key is pressed. It will remove the intro and
-- start the simulation.
--------------------------------------------------------------------------------
function love.keypressed(key, isrepeat)
  show_intro = false
end


--------------------------------------------------------------------------------
-- Callback that is called after each update. It simple draws the start line, 
-- the balls and the floor.
--------------------------------------------------------------------------------
function love.draw()
  -- Push the view and world matrix.
  love.graphics.push()

  -- Hack to determine if the screen uses high DPI, in which case we scale 
  -- everything up.
  local is_high_dpi = (love.graphics.getHeight() / 600) == 2;
  if is_high_dpi then
    love.graphics.scale(2, 2)
  end

  -- Draw the start line.
  love.graphics.rectangle('fill', 0, START_Y + BALL_RADIUS, 800, 3)

  -- Draw the balls.
  for i,v in pairs(balls) do
    local title = v.interpolation .. ' with ' .. v.sim_step 
      .. ' step(s) per frame'
    love.graphics.printf(title, v.pos.x - 50, 10, 100, 'center')
    love.graphics.circle('fill', v.pos.x, v.pos.y, BALL_RADIUS, 100)
  end

  -- Draw the floor.
  love.graphics.rectangle('fill', 0, FLOOR_POS_Y, 800, 10)

  if show_intro then
    local intro_message = 'Press space or any other key to start the simulation'
    love.graphics.printf(intro_message, 400-150, 250, 300, 'center')
  end

  -- Pop the scaler.
  love.graphics.pop()
end


--------------------------------------------------------------------------------
-- Vec2
-- =============================================================================
-- 
-- The Vec2 (short for Vector2) is a simple utility class for position and 
-- velocity. The primary reason is to make vector math a bit easier to read.
--
--------------------------------------------------------------------------------
Vec2 = {}
Vec2.__index = Vec2


--------------------------------------------------------------------------------
-- Constructor of Vec2, which has two parameters. If one of the parameters is 
-- omitted then it is assumed to be `0`.
--------------------------------------------------------------------------------
function Vec2.new(x, y)
  local default_table = { x = x or 0, y = y or 0 }
  return setmetatable(default_table, Vec2)
end


--------------------------------------------------------------------------------
-- The add operator adds the individual componets of two vectors together and 
-- returns this is a new Vec2.
--------------------------------------------------------------------------------
function Vec2.__add( v1, v2 )
  return Vec2.new(v1.x + v2.x, v1.y + v2.y)
end


--------------------------------------------------------------------------------
-- The multiply operator multiplies the individual componets of two vectors and 
-- returns this is a new Vec2. If one of the parameters is number then this 
-- number is used to multiply both components of the vector. 
--------------------------------------------------------------------------------
function Vec2.__mul(v1, v2)
  if type(v1) == 'number' then v1 = Vec2.new(v1, v1) end
  if type(v2) == 'number' then v2 = Vec2.new(v2, v2) end
  return Vec2.new(v1.x * v2.x, v1.y * v2.y)
end