pendulum/distributed/worker.py

58 lines
1.6 KiB
Python

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 <pavle.portic@tilda.center>
#
# Distributed under terms of the BSD 3-Clause license.
import numpy as np
from scipy.integrate import odeint
from .app import app
from .common import iterate_theta
g = 9.81
def deriv(y, t, L1, L2, m1, m2):
theta1, z1, theta2, z2 = y
c, s = np.cos(theta1 - theta2), np.sin(theta1 - theta2)
theta1dot = z1
z1dot = (m2 * g * np.sin(theta2) * c - m2 * s * (L1 * z1**2 * c + L2 * z2**2) - (m1 + m2) * g * np.sin(theta1)) / L1 / (m1 + m2 * s**2)
theta2dot = z2
z2dot = ((m1 + m2) * (L1 * z1**2 * s - g * np.sin(theta2) + g * np.sin(theta1) * c) + m2 * L2 * z2**2 * s * c) / L2 / (m1 + m2 * s**2)
return theta1dot, z1dot, theta2dot, z2dot
def solve(L1, L2, m1, m2, tmax, dt, y0):
t = np.arange(0, tmax + dt, dt)
y = odeint(deriv, y0, t, args=(L1, L2, m1, m2))
return y[:, 0], y[:, 2]
def compute(L1, L2, m1, m2, tmax, dt, theta1_init, theta2_init):
y0 = np.array([theta1_init, 0, theta2_init, 0])
theta1, theta2 = solve(L1, L2, m1, m2, tmax, dt, y0)
return theta2_init, theta1[-1], theta2[-1]
@app.task
def compute_single(L1, L2, m1, m2, tmax, dt, theta1_init, theta2_init):
return compute(L1, L2, m1, m2, tmax, dt, theta1_init, theta2_init)
@app.task
def compute_chunk(L1, L2, m1, m2, tmax, dt, theta1_init, theta_resolution):
return (
theta1_init,
[
compute(L1, L2, m1, m2, tmax, dt, theta1_init, theta2_init)
for theta2_init in iterate_theta(theta_resolution)
]
)