# Visual solution of the Newtonial differential equation with Python

I wrote this Python code that plots the movement of an object under the effect of a given force function in 2D by solving the Newton's movement equation numerically. One can add other force functions, or even parameters to the draw_path function. I tried to make it as readable as I could. I would really appreciate if you could tell me what did I wrong, and what would have you done diferently.

Since I learned programming only by tutorials and codes, never from proper lessons/courses, and this is my first finished code, I probably did some weird things. Since I am not a native speaker, I probably wrote some weird comments. Sorry.

#!/usr/bin/env python

#  movement_equation_2.0.py
#
#
#  This program is free software; you can redistribute it and/or modify
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  :Author: Nagy Gergely
#  :Version: 0.2.1 beta
#  :Status: Prototype
#  :Date: 2016.01.05

"Non-physics functions"

from PIL import Image, ImageDraw

#dimensions for vectors
x = 0
y = 1

def plot_base(xmin, xmax, ymin, ymax):
"""
Creates a base for the plot: a white picture with two orthogonal lines, the x
and y axis, according to the given minimum and maximum coordinates in pixels.
"""
base_color = 'white'
axis_color = 'grey'

xsize = abs(xmax - xmin)
ysize = abs(ymax - ymin)
plot = Image.new('RGB', (xsize, ysize), base_color)
draw = ImageDraw.Draw(plot)

#draw x axis if shown:
if ymin < 0 and ymax > 0:
draw.line((0, ymax, xsize, ymax), axis_color)
elif ymax > 0:
draw.line((0, ysize, xsize, ysize), axis_color)
else:
draw.line((0, 0, xsize, 0), axis_color)

#draw y axis if shown:
if xmin < 0 and xmax > 0:
draw.line((0-xmin, 0, 0-xmin, ysize), axis_color)
elif xmax > 0:
draw.line((0, 0, 0, ysize),axis_color)
else:
draw.line((xsize, 0, xsize, ysize),axis_color)

return plot

def draw_path(m, r0, v0, force_function, time, dt, plot_size, resolution):
"""
Draws the path of an object with weight m starting from r0 with velocity v0,
according to the force function given in the force_function method.

:param m: the mass of the object
:param r0: the coordinates of the object at t=0 in meters (2-tuple)
:param v0: the coordinates of the initial velocity vector (2-tuple)
:param force_function: the function of the force applied to the object, described below
:param time: the time of the movement
:param dt: time elapsed ed between two calculated point: the bigger it is, the faster but less accurate is the result
:param plot_size: the minimum and the maximum x and y coordinates as a 4-tuple: (xmin, xmax, ymin, ymax)
:param resolution: resolution of the result plot image (meter/pixel)
"""
trace_color = (255, 0, 0)
plot_pixsize = [int(i/resolution) for i in plot_size]
plot = plot_base(*plot_pixsize)
draw = ImageDraw.Draw(plot)
m = float(m)
r = (float(r0[x]), float(r0[y]))
v = (float(v0[x]), float(v0[y]))

i = 0
pos, prev_pos = None, None
while float(i)*dt < time:

#physics calculations
F = force_function(m=m, r=r, v=v)
a = (F[x]/m                        , F[y]/m                       )
v = (v[x] + a[x]*dt                , v[y] + a[y]*dt               )
r = (r[x] + v[x]*dt + a[x]/2*dt**2 , r[y] + v[y]*dt + a[y]/2*dt**2)

#drawing the line between the current and the previous position if both is on the plot
pix_pos_x = round(r[x]/resolution) - plot_pixsize
pix_pos_y = round(r[y]/resolution) *-1 + plot_pixsize
if ((0 < pix_pos_x < plot_pixsize-plot_pixsize) and
(0 < pix_pos_y < plot_pixsize-plot_pixsize)):
pos = pix_pos_x, pix_pos_y
else:
pos = None

if pos is not None and prev_pos is not None:
draw.line((prev_pos, pos), fill=trace_color)
#plot.putpixel((pix_pos_x, pix_pos_y), trace_color)     #'dotty' but may visualize speed

prev_pos=pos
i += 1

return plot

"""
*******************************************************************************
FORCE FUNCTIONS
*******************************************************************************
"""

import math

"physics constants"

_c_  = 299792458 # m/s                  speed of light
_y_  = 6.67384 * (10**-11) # Nm²/kg²    gravitational constant
_h_  = 6.62606957 * (10**-34) # Js      Planck's constant
_E0_ = 8.854187817 * (10**-12) # C²/Nm² electric constant
_u0_ = 4.0 * math.pi # Tm/A             magnetic constant
_g_  = 9.80665 # m/s²                   standard gravity
_e_  = 1.602176565 * (10**-19) # C      elementary charge
_me_ = 9.10938291 * 10**-31 # kg        mass of electron
_mp_ = 1.672621777 * 10**-27 #kg        mass of proton
_Na_ = 6.02214129 * 10**23 #1/mol       Avogadro's constant

_Me_  = 5.972 * 10**24 # kg             mass of the Earth
_Ms_  = 1.989 * 10**30 # kg             mass of the Sun
_Mm_  = 7.34767309 * 10**22 #kg         mass of the Moon

"""
Force functions for the movement plotter. They should take the paramaters as keyword
arguments (use **kwargs to be compatible with more parameters in the future), and
return the coordinates of the force vector at these parameters as a 2-tuple.
Possible variablesat the moment: m, r, v
"""

def r_xy_dep(**variables):
"""force depends on the x and the y coordinates"""
r = variables['r']
F = [0.0 , -1.0]

F[x] = 0
F[y] = _g_*F[y]

return F

def v_xy_dep(**variables):
"""force depends on the x and the y velocity"""
v = variables('v')
F = [0.0 , 1.0]
F[x] = -(v[x]**2) * v[x]/abs(v[x])
F[y] = -(v[y]**2) * v[y]/abs(v[y])

return F

def central(**variables):
"""central force field, force depends on the
vector from the centrum to the object"""

r0 = variables['r']
c = (0.0 , 0.0)               # centrum
r = (r0[x]-c[x] , r0[y]-c[y]) # vector from centrum to r0
r_ = math.hypot(r[x], r[y])   # length of r
fi = math.atan2(r[y], r[x])   # angle of r

F_ = -r_                      #force dependency
dfi = 0                       #the angle between the force vector and r

F_x = F_*math.cos(fi+dfi)
F_y = F_*math.sin(fi+dfi)
return (F_x, F_y)

def gravitational(**variables):
#(special type of central dependency)

central_mass = _Me_   # mass of the scource object
m = variables['m']
r0 = variables['r']
c = (0.0 , 0.0)
r = (r0[x]-c[x] , r0[y]-c[y])
r_ = math.hypot(r[x], r[y])
fi = math.atan2(r[y], r[x])
F_ = -r_
F_x = F_*math.cos(fi+dfi)
F_y = F_*math.sin(fi+dfi)
return (F_x, F_y)

def v_dep(**variables):
"""force depends on the velocity vector
(i.e. charged particle in magnetic field)"""
v_xy = variables['v']
v_ = math.hypot(v_xy[x], v_xy[y])
fi = math.atan2(v_xy[y], v_xy[x])

F_ = v_
dfi = math.pi/2

F_x = F_*math.cos(fi+dfi)
F_y = F_*math.sin(fi+dfi)
return (F_x, F_y)

"""
******************************************************************************
MAIN MOVEMENT DRAWER FUNCTION
******************************************************************************
"""

def main(arg):

m = 0.1               #mass of the object in kilograms
r0 = (0, 0)           #initial coordinates in meters
v0 = (30, 0)          #initial velocity vector in m/s
dependency = v_dep    #force function from above
time = 10             #time of movement to draw in secs
time_res = 1/70000    #time between steps in secs, determines accuracy and running time
plot_size = (-10, 10, -10, 10)  #minimum and maximum coordinates in meters (xmin, xmax, ymin, ymax)
plot_res = 0.01       #meters per pixel on the plot

plot = draw_path(m, r0, v0, dependency, time, time_res, plot_size, plot_res)
plot.show()

if __name__ == '__main__':
import sys
sys.exit(main(sys.argv))


Also on Github: https://gist.github.com/godot11/998e71fca8f8f4fce1a1

This is pretty good work for a self-taught programmer. I'd like to make a few suggestions to improve readability.

First, I would define a 2D vector class. That would allow you to simplify expressions where you are doing the same operation to both the x and the y coordinate, such as r = (r[x] + v[x]*dt + a[x]/2*dt**2 , r[y] + v[y]*dt + a[y]/2*dt**2), into

r += (v * dt) + (a * dt**2 / 2)


To achieve that, I would use namedtuple, so that you can have vec.x and vec.y members instead of vec[x] and vec[y].

from collections import namedtuple
import math
from PIL import Image, ImageDraw

class Vec2(namedtuple('Vec2', 'x y')):
return Vec2(self.x + other.x, self.y + other.y)

def __sub__(self, other):
return Vec2(self.x - other.x, self.y - other.y)

def __mul__(self, scale):
return Vec2(self.x * scale, self.y * scale)

def __truediv__(self, scale):
return Vec2(self.x / scale, self.y / scale)

def rotate(self, angle):
sin, cos = math.sin(angle), math.cos(angle)
return Vec2(self.x * cos - self.y * sin, self.x * sin + self.y * cos)


The Vec2 class would also let the code express your mathematical intentions — for example, your entire eight-line v_dep() function could be written as…

###############################################################################
# FORCE FUNCTIONS
###############################################################################
# Force functions for the movement plotter. They should take the paramaters as
# keyword arguments (use **kwargs to be compatible with more parameters in the
# future), and return the coordinates of the force vector at these parameters
# as a 2-tuple.  Possible variables at the moment: m, r, v
###############################################################################

def v_dep(**variables):
"""force depends on the velocity vector
(i.e. charged particle in magnetic field)"""
return variables['v'].rotate(math.pi / 2)


The other change would be to disentangle the physics from the plotting. In particular, I'd prefer not to have code like pix_pos_x = round(r[x]/resolution) - plot_pixsize in your draw_path loop, which is already complicated as it is.

I suggest defining a Canvas class to contain the image, the bounds information, and handle the physical-to-pixel coordinate translation.

class Canvas:
def __init__(self, bounds, resolution, base_color='white'):
"""
:param bounds: positions of the upper-left and lower-right corners, each as a Vec2
:param resolution: resolution of the result plot image (meter/pixel)
"""
self._nw_corner, self._se_corner = bounds
self._resolution = resolution

width = round((self._se_corner.x - self._nw_corner.x) / resolution)
height = round((self._nw_corner.y - self._se_corner.y) / resolution)
self._plot = Image.new('RGB', (width, height), base_color)
self._draw = ImageDraw.Draw(self._plot)

self._pos = None
self.color = 'black'

def draw_axes(self, axis_color='grey'):
#draw x axis
self._draw.line([
self._phys_to_draw(Vec2(self._nw_corner.x, 0)),
self._phys_to_draw(Vec2(self._se_corner.x, 0))
], axis_color)

#draw y axis
self._draw.line([
self._phys_to_draw(Vec2(0, self._nw_corner.y)),
self._phys_to_draw(Vec2(0, self._se_corner.y))
], axis_color)

def draw_to(self, r):
"""Draw a line from the previous position (if any) to the specified position"""
self._prev_pos, self._pos = self._pos, self._phys_to_draw(r)
if self._prev_pos is not None:
self._draw.line((self._prev_pos, self._pos), fill=self.color)

def show(self):
self._plot.show()

def _phys_to_draw(self, r):
"""Translate physical coordinates to image coordinates"""
return (
round((r.x - self._nw_corner.x) / self._resolution),
round((self._nw_corner.y - r.y) / self._resolution)
)


Here's the rest of the code:

def draw_path(canvas, m, r0, v0, force_function, duration, dt):
"""
Draws the path of an object with weight m starting from r0 with velocity v0,
according to the force function given in the force_function method.

:param m: the mass of the object
:param r0: the coordinates of the object at t=0 in meters (2-tuple)
:param v0: the coordinates of the initial velocity vector (2-tuple)
:param force_function: the function of the force applied to the object, described below
:param duration: the duration of the movement
:param dt: time elapsed between two calculated points: the bigger it is, the faster but less accurate is the result
"""
r = r0
v = v0

canvas.color = (255, 0, 0)
canvas.draw_to(r)
for _ in range(int(duration / dt)):
F = force_function(m=m, r=r, v=v)
a = F / m
v += a * dt
r += (v * dt) + (a * dt**2 / 2)
canvas.draw_to(r)

###############################################################################
# MAIN MOVEMENT DRAWER FUNCTION
###############################################################################

def main(arg):
canvas = Canvas(
bounds=(Vec2(-10, 10), Vec2(10, -10)),
resolution=0.01,                #meters per pixel on the plot
)
canvas.draw_axes()
draw_path(
canvas=canvas,
m=0.1,                          #mass of the object in kilograms
r0=Vec2(0, 0),                  #initial coordinates in meters
v0=Vec2(30, 0),                 #initial velocity vector in m/s
force_function=v_dep,           #force function from above
duration=10,                    #time of movement to draw in secs
dt=1/70000,                     #time between steps in secs, determines accuracy and running time
)
canvas.show()

if __name__ == '__main__':
import sys
sys.exit(main(sys.argv))


• Avoid using """strings""" as comments. Either write proper """docstrings""" on functions, or write # comments.
• Numbers like 6.67384 * (10**-11) are better written as 6.67384e-11.
• Counting loops like…

i = 0
while …:
…
i += 1


… are better written using for i in range(…): ….

• By moving the creation of the image and axes out of draw_path(), we can allow for multiple particles to be drawn on the same canvas.
• draw_path() takes a lot of parameters. Using named parameters makes the code clearer and less susceptible to parameter mismatches.
• Don't worry about drawing out of bounds. Let the PIL library crop the image for you.
• Thank you very much for the review! You clearly spent time with it, and I learned a lot of things from it. Jan 5, 2016 at 22:37
• There's a simpler way to implement rotate, using only two trig calls: x, y = self; s, c = sin(theta), cos(theta); return Vec2(x * c - y * s, x * s + y * c) Jan 8, 2016 at 21:37
• @GarethRees Incorporated in Rev 3. Thanks. Jan 8, 2016 at 22:28