# Exercise 3: Using types - a simple N-body in 80 lines

Julia makes it easy to write your own types, which are similar to structs in C-family languages. Multiple dispatch adds a powerful way to handle tasks like default values or different behaviours really easily.

Here we'll look at how you could construct a simple N-body code. We'll assume everything has mass 1 to keep it simple.

• define an abstract type representing the bodies in the simulation
• define subtypes of this representing two different kinds of bodies in 2D: a moving body and a fixed body
• define a type which holds all the bodies in the simulation
• make a few of these bodies in various positions to start
• write a movement function to move the bodies as necessary
• write a force function to give an inverse square law force between bodies
• write an acceleration function to apply gravity to the bodies
• write a step function to bring it all together (acceleration and movement)

Extras:

• write plot functions for individual points and for the whole system
• plot it and watch the carnage

## Define an abstract type representing the bodies in the simulation

This code block demonstrates how to define an abstract type: a type whose only purpose is to encompass other types.

``````

In [1]:

abstract Body

``````

## Define subtypes of this representing moving and fixed bodies in 2D

This code block demonstrates how to define "concrete types" which actually contain values. We declare two types, MovingBody and FixedBody, and make them subtypes of Body (by writing `<: Body`). We also write a function `MovingBody(x,y)` which makes a MovingBody that defaults to zero movement.

``````

In [2]:

type MovingBody <: Body
x::Float64
y::Float64
dx::Float64
dy::Float64
end
MovingBody(x, y) = MovingBody(x, y, 0, 0)

type FixedBody <: Body
x::Float64
y::Float64
end

``````

## Define a type which holds all the bodies in the simulation

This code block demonstrates the syntex to use when specifying the types of arrays. We could have written `Array{Body, 1}` to specify a 1D array containing `Body` types, but `Vector{Body}` is a slightly shorter way of writing the same thing.

``````

In [3]:

type System
bodies::Vector{Body}
end

``````

## Make a few of these in various positions to start

This code block demonstrates how we can make instances of the body types we just defined. It's nice to put as much of this in a function as possible.

``````

In [4]:

function setup()
p1 = FixedBody(0,0)
p2 = MovingBody(1,0) # defaults to no motion ... like writing MovingBody(1,0,0,0)
p3 = MovingBody(1,1,1,1)
p4 = MovingBody(0,1,1,-1)
bodies = [p1, p2, p3, p4]

System(bodies)
end

sys = setup()

``````
``````

Out[4]:

System(Body[FixedBody(0.0,0.0),MovingBody(1.0,0.0,0.0,0.0),MovingBody(1.0,1.0,1.0,1.0),MovingBody(0.0,1.0,1.0,-1.0)])

``````

## Write a movement function to move the bodies as necessary

The convention in Julia is to put a ! on the end of the function name if it modifies its arguments. Since we're actually changing the body here (we alter the body's x and y positions), we call the function `move!`.

``````

In [5]:

const dt = 0.1  # declare it as const because it's not in a function

function move!(b::MovingBody)
b.x += b.dx * dt
b.y += b.dy * dt

b # we return the body here - its coordinates have now been altered
end

move!(b::FixedBody) = b

``````
``````

Out[5]:

move! (generic function with 2 methods)

``````

See how we just wrote a second version of `move!` but defined it to only be applicable to a FixedBody? Now we can call `move!` on either a MovingBody or a FixedBody and know that it will behave appropriately.

## Write a force function to give an inverse square law force between bodies

This code block shows how you can use Unicode characters like ² (type `\^2` and press Tab) in variable names to make them more compehensible.

``````

In [6]:

function force(b1::Body, b2::Body)
dx = b2.x - b1.x
dy = b2.y - b1.y
d² = dx^2 + dy^2 + 1e-9  # we add a 1e-9 softening term to avoid dividing by zero
inv_d² = 1 / d²
Fx = dx * inv_d²
Fy = dy * inv_d²

Fx, Fy
end

``````
``````

Out[6]:

force (generic function with 1 method)

``````

## Write an acceleration function to apply gravity to the bodies

This function needs to include the system (group of bodies) as an argument, because each body is affected by the gravity of all other bodies in the system. Another option would be to make `sys`, the `System`, a global variable. But this option is nicer style and lets us make multiple `System`s if we want.

``````

In [7]:

function accelerate!(sys::System, b::MovingBody)
d2x = 0.0
d2y = 0.0
for other in sys.bodies
if b !== other
Fx, Fy = force(b, other)
d2x += Fx
d2y += Fy
end
end
b.dx += d2x * dt
b.dy += d2y * dt

b
end

accelerate!(sys::System, b::FixedBody) = b

``````
``````

Out[7]:

accelerate! (generic function with 2 methods)

``````

## Write a step function to bring it all together

By this stage, the step function is really simple to write. We write two versions: one which steps an individual particle, and one which calls the same function but on every particle in a system.

``````

In [8]:

function step!(sys::System, b::Body)
accelerate!(sys, b)
move!(b)
end

function step!(sys::System)
for b in sys.bodies
step!(sys, b)
end
end

``````
``````

Out[8]:

step! (generic function with 2 methods)

``````

## Extra: plot functions to plot individual points and the system

Since the `plot` and `scatter` (etc) commands from Matplotlib overlay plots by default, we can take advantage of this by writing different `plot` commands for different points. We make fixed things red, moving things black with direction arrows, and then write a third `plot` command to plot everything in the system by just calling `plot` on each individual point.

``````

In [9]:

using PyPlot

function plot(b::FixedBody)
scatter(b.x, b.y, s=50, color="Red")
end

function plot(b::MovingBody)
scatter(b.x, b.y, s=50, color="Black")
arrow(b.x, b.y, b.dx * dt/2, b.dy * dt/2, facecolor="Black")
end

function plot(s::System)
for b in s.bodies
plot(b)
end
axis("equal")
end

``````
``````

Out[9]:

plot (generic function with 3 methods)

``````
``````

In [11]:

sys = setup()
plot(sys)

``````
``````

Out[11]:

(-0.2,1.2000000000000002,-0.2,1.2000000000000002)

``````

## Watch the carnage

Here's what that looks like if we plot the system over 5 time units. By overplotting, we build up a simple trajectory.

``````

In [12]:

for t = 0:dt:5
plot(sys)
step!(sys)
end

``````
``````

``````