Model of kinematics and energy

PhysicsModel.py

import cv2
import numpy as np

import math

# standard units: time = sec, displacement = pixels, mass = kg, force = N

conv = 100
g = 10 * conv
slope_int = 0.01
fps = 1000

def ir(num):
    return int(round(num))

class TerrainModel:
    def __init__(self):
        self.functions = [] # form ((min, max), lambda)
        self.length = 0

    def add(self, minmax, function):
        self.functions.append((minmax, function))
        self.length += 1

    def f(self, x):
        for ind in range(0, self.length):
            func = self.functions[ind]
            if x >= func[0][0] and x <= func[0][1]:
                return func[1](x)
        return None

class Ball:
    def __init__(self, radius, mass, initpos=(0, 0), initvel=(0, 0), rendersize=None):
        self.radius = radius
        self.mass = mass

        self.pos = list(initpos)
        self.vel = list(initvel)

        if rendersize is not None:
            self.rendersize = rendersize
        else:
            self.rendersize = self.radius

    def move(self, t=1.0, a=(0, 0)):
        self.pos[0] += self.vel[0] * t
        self.pos[1] += self.vel[1] * t

        self.vel[0] += a[0] * t
        self.vel[1] += a[1] * t


    def checkIfWithin(self, point):
        if point[0] != None and point[1] != None:
            if (point[0] - self.pos[0]) ** 2 + (point[1] - self.pos[1]) ** 2 < self.radius ** 2:
                return True
        return False

    def calcEnergy(self):
        ke = 0.5 * self.mass * ((self.vel[0] / conv) ** 2 + (self.vel[1] / conv) ** 2)
        ge = self.mass * (g / conv) * (height - self.pos[0]) / conv

        return {"ke" : ke, "ge" : ge}

    def render(self, img):
        tmp = img.copy()
        tmp = cv2.circle(tmp, (ir(self.pos[1]), ir(self.pos[0])), radius=self.rendersize, color=(0, 255, 0), thickness=-1)

        return tmp

class Spring:
    def __init__(self, k, start, length, func, angle, dir=True): # dir: True start left expands forward, False start right expand backwards
        self.k = k

        self.start = start
        self.length = length
        self.compress = 0
        self.func = func
        self.angle = angle
        self.dir = dir

        self.cip = 20

    def update(self, obj):
        self.compress = 0

        ra = range(self.start[1] * self.cip, ir(self.start[1] + self.length * np.cos(self.angle)) * self.cip)
        if not self.dir:
            ra = range(ir(self.start[1] - self.length * np.cos(self.angle)) * self.cip, self.start[1] * self.cip)
        for self.xcip in ra:
            x = self.xcip/self.cip
            if obj.checkIfWithin((self.func(x), x)):
                endleny = self.func(self.length * np.cos(self.angle)) - self.func(x)
                endlenx = self.start[1] + self.length * np.cos(self.angle) - x

                if not self.dir:
                    endleny = self.start[0] - self.func(self.length * np.cos(self.angle)) - self.func(x)
                    endlenx = self.start[1] - self.length * np.cos(self.angle) - x

                self.compress = math.sqrt(endleny ** 2 + endlenx ** 2)

                if self.dir:
                    return

    def calcEnergy(self):
        el = 0.5 * self.k * ((self.compress / conv) ** 2)

        return {"el" : el}

    def getDAccel(self, m):
        a = (self.k * self.compress) / m

        if not self.dir:
            a *= -1

        return (np.sin(self.angle) * a, np.cos(self.angle) * a)

    def render(self, img):
        tmp = img.copy()
        endx = abs(self.start[1] - (self.length - self.compress) * np.cos(self.angle))
        tmp = cv2.line(tmp, (ir(self.start[1]), ir(self.start[0])),
                       (ir(endx), ir(self.func(endx))), color=(0, 0, 255))

        return tmp


obj = Ball(0.1, 20, initpos=(0, 0), rendersize=5)

spr = Spring(200, (ir(6 * conv), ir(10.4 * conv)), 6 * conv, lambda x : 0.57735 * x, np.arctan(0.57735), dir=False)

height = 600
width = 1040

model = TerrainModel()
model.add((0, 1038), lambda x : 0.57735 * x)

canvas = np.zeros([height, width, 3], dtype=np.uint8)
canvas.fill(255)

for x in range(0, width):
    y = model.f(x)
    if y != None:
        if ir(y) < height and ir(y) >= 0 and x < width and x >= 0:
            canvas[ir(y)][x] = (0, 0, 0)
        else:
            print("Out of bounds", (ir(y), x))

time = 0
num = 0

cv2.startWindowThread()
cv2.namedWindow("img")

f = open("log.txt", "w")

while True:
    a = (0, 0)
    if not obj.checkIfWithin((model.f(obj.pos[1]), obj.pos[1])):
        a = (g, 0)
    else:
        back_y = model.f(obj.pos[1] - slope_int)
        for_y = model.f(obj.pos[1] + slope_int)

        slope = 0

        if back_y != None and for_y != None:
            slope = (for_y - back_y) / (2 * slope_int)

        elif back_y == None:
            slope = (for_y - model.f(obj.pos[1])) / (slope_int)

        elif for_y == None:
            slope = (model.f(obj.pos[1]) - back_y) / (slope_int)

        if slope == 0:
            obj.vel[0] = 0
        else:
            theta = np.arctan(slope)
            a = (g * np.sin(theta) * np.sin(theta), g * np.sin(theta) * np.cos(theta))

    spr.update(obj)

    if spr.compress != 0:
        daccel = spr.getDAccel(obj.mass)
        a = (a[0] + daccel[0], a[1] + daccel[1])

    obj.move(t=1.0/fps, a=a)

    energy = {}
    energy.update(obj.calcEnergy())
    energy.update(spr.calcEnergy())

    img = canvas.copy()
    img = obj.render(img)
    img = spr.render(img)

    cv2.imshow("img", img)
    #cv2.waitKey(1)
    time += 1.0/fps

    if ir(time * 10000) / 100.0 % 1 == 0:
        cv2.waitKey(1)

        print(round(time, 2), [(ele, ir(val)) for ele, val in energy.items()], ir(sum(energy.values())))
        print(obj.pos, spr.compress)
        num += 1
        cv2.imwrite("img/frame" + str(num) + ".jpg", img)
        f.write('%.2f\t%d\t%d\t%d\t%d\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f\n' %
                (round(time, 2),
                 ir(energy["ke"]), ir(energy["ge"]), ir(energy["el"]), ir(sum(energy.values())),
                 round(obj.pos[1], 2), round(obj.pos[0], 2),
                 round(obj.vel[1], 2), round(obj.vel[0], 2),
                 round(a[1], 2), round(a[0], 2),
                 round(spr.compress, 2)))