7
\$\begingroup\$

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
#
#  Copyright 2016 Nagy Gergely
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  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[0]
        pix_pos_y = round(r[y]/resolution) *-1 + plot_pixsize[3]
        if ((0 < pix_pos_x < plot_pixsize[1]-plot_pixsize[0]) and
            (0 < pix_pos_y < plot_pixsize[3]-plot_pixsize[2])):
            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

\$\endgroup\$
11
\$\begingroup\$

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')):
    def __add__(self, other):
        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[0] 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))

Additional remarks:

  • 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.
| improve this answer | |
\$\endgroup\$
  • \$\begingroup\$ Thank you very much for the review! You clearly spent time with it, and I learned a lot of things from it. \$\endgroup\$ – Neinstein Jan 5 '16 at 22:37
  • 1
    \$\begingroup\$ 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) \$\endgroup\$ – Gareth Rees Jan 8 '16 at 21:37
  • 1
    \$\begingroup\$ @GarethRees Incorporated in Rev 3. Thanks. \$\endgroup\$ – 200_success Jan 8 '16 at 22:28

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.