class TestModel(unittest.TestCase):
def setUp(self):
# test data
self.circle_solver = SystemSolver(CircleSystem())
def testSomething(self):
res = self.circle_solver.run([0, 7], [0, 1])['results']
self.assertTrue(res.success)
for i, t in enumerate(res.t):
x, y = res.y[:, i]
x_exp, y_exp = np.array([-np.sin(t), np.cos(t)])
self.assertAlmostEqual(x, x_exp, places=2)
self.assertAlmostEqual(y, y_exp, places=2)
def testCircleBox(self):
t_span = [0, 300]
res = self.circle_solver.run(t_span, [0, 1])
fig = self.circle_solver.plot2d(res, None)
ax = fig.get_axes()[0]
t_dense = np.arange(*t_span, t_span[1]/1000)
box = np.array([np.cos(t_dense), -np.sin(t_dense)])
ax.plot(*box, '-')
newbox = np.sqrt(res['results'].y[0, -1]**2 +
res['results'].y[1, -1]**2)*box
ax.plot(*newbox, '-')
plt.show(False)
class TestModel(unittest.TestCase):
def setUp(self):
# test data
self.pendulum_system = DoublePendulumSystem()
self.pendulum_solver = SystemSolver(self.pendulum_system)
self.run_dp_render = False
def testDoublePendulum(self):
clearFigs()
run = self.pendulum_solver.run([0, 2], [.2, 1, 0, 0])
run2 = self.pendulum_solver.run([0, 2], [.2, 1.1, 0, 0])
fig = self.pendulum_solver.plotnd(run)
self.pendulum_solver.plotnd(run2, fig)
# plt.show(True)
clearFigs()
if self.run_dp_render:
system = DoublePendulumSystem()
fig = system.render_path(run, dot_size=2)
system.render_path(run2, fig=fig, dot_size=2)
plt.show(True)
def testDoublePendulumFade(self):
clearFigs()
run = self.pendulum_solver.run([0, 5], [.2, 1, 0, 0])
system = DoublePendulumSystem()
if self.run_dp_render:
fig = system.render_fade_trail(run)
plt.show(False)
step = .02
for t in np.arange(.5, 5, step)[1:]:
plt.pause(.001)
plt.clf()
fig = system.render_fade_trail(run, fig=fig, time=t)
def testDoublePendulumTrail(self):
clearFigs()
run = self.pendulum_solver.run([0, 10], [.2, 1, 0, 0])
system = DoublePendulumSystem()
if self.run_dp_render:
fig = system.render_trail(run, time=.5)
plt.show(False)
step = .05
for t in np.arange(.5, 10, step):
plt.pause(.00001)
plt.clf()
fig = system.render_trail(run, fig=fig, time=t)
def testDoublePendulumTrailFast(self):
clearFigs()
run = self.pendulum_solver.run([0, 10], [.2, 1, 0, 3])
system = DoublePendulumSystem()
if self.run_dp_render:
fig = system.render_trail(run, time=.5)
plt.show(False)
step = .05
for t in np.arange(.5, 10, step):
plt.pause(.00001)
plt.clf()
fig = system.render_trail(run, fig=fig, time=t)
def testWeirdLenSys(self):
sys = DoublePendulumSystem(1, 2, .5, 1.5)
slv = SystemSolver(sys)
tf = 3
run1 = slv.run([0, tf], [pi, 1, 1, 0])
run2 = slv.run([0, tf], [pi, 1.001, 1, 0])
if self.render:
self.runRender([run1, run2], 'odd_pend')
def testOrbit(self):
Ms = 1.98847 * 10**30 # kg
# M_earth = 5.9722 * 10**24 # kg
Mj = 1.899 * 10**27 # kg
Rj = 778.3 * 10**9 # m
Tj = 3.743 * 10**8 # s
# V_earth = 30 * 10**3 # m/s
# r = 149.6 * 10**9 # m
vel = 2 * pi * Rj / Tj
sun_x = np.zeros(3)
sun_v = np.zeros(3)
jup_x = np.zeros(3)
jup_v = np.zeros(3)
jup_x[0] = Rj
jup_v[1] = vel
m = [Ms, Mj]
v = np.concatenate((sun_x, jup_x, sun_v, jup_v)).reshape(-1)
sys = NBodySystem(m)
solver = SystemSolver(sys)
res = solver.run([0, Tj * 1.5], v)
solver.plotnd(res)
def testTwoPoints(self):
# using IC from TODO
sys = NBodySystem(body_masses=[1, 1], G=1)
solver = SystemSolver(sys)
max_t = 52
tspan = [0, max_t]
y0 = np.zeros(2 * sys.body_dim * len(sys.body_masses),
dtype=np.float64)
y0[0:3] = [1, 0, 0]
y0[3:6] = [-1, 0, 0]
y0[6:9] = [0, .1, 0]
y0[9:12] = [0, -.1, 0]
run = solver.run(tspan, y0)
# print()
# print(*y0)
# print(*run['results'].y[:, -1])
# clearFigs()
# fig = solver.plotnd(run)
# for axes in fig.axes:
# axes.plot(np.zeros_like(run['results'].y[0]))
t = run['results'].t
y = run['results'].y
# print(t.shape, y.shape)
# print(t[-1], max_t)
y0 = y[0]
y1 = y[3]
clearFigs()
plt.figure()
plt.plot(t, y0)
plt.plot(t, y1)
def testFig8(self):
# using IC from TODO
sys = NBodySystem(body_masses=[1, 1, 1], G=1)
solver = SystemSolver(sys)
tspan = [0, 10]
y0 = np.zeros(2 * sys.body_dim * len(sys.body_masses),
dtype=np.float64)
x1 = np.array([0.97000436, -0.24308753, 0])
x3p = np.array([-0.93240737, -0.86473146, 0])
y0[0:3] = x1
y0[3:6] = -x1
y0[6:9] = 0
y0[9:12] = -x3p / 2
y0[12:15] = -x3p / 2
y0[15:18] = x3p
# print(sys.fun(np.zeros_like(y0), y0).reshape(6, -1))
run = solver.run(tspan, y0)
clearFigs()
solver.plotnd(run)
# print(run['results'].y[:, -1].reshape(6, -1))
y_act = run['results'].y[:9]
run['results'].y = y_act[:3]
fig = solver.plot3d(run)
run['results'].y = y_act[3:6]
fig = solver.plot3d(run, fig=fig)
run['results'].y = y_act[6:9]
fig = solver.plot3d(run, fig=fig)
def testRandomInit(self):
sys = NBodySystem(body_masses=[1, 1, 1], G=1)
solver = SystemSolver(sys)
tspan = [0, 100]
expand = 10
y0 = np.random.rand(2 * sys.body_dim * len(sys.body_masses)) * expand
y0[9:] /= expand / 2
total_mass = np.sum(sys.body_masses)
vcm = np.zeros(3)
for m, v in zip(sys.body_masses, y0[9:].reshape(3, -1)):
vcm += m * v
vcm /= total_mass
yps = y0[9:].reshape(3, -1)
yps -= vcm[None, :]
run = solver.run(tspan, y0)
clearFigs()
solver.plotnd(run)
# print(run['results'].y[:, -1].reshape(6, -1))
y_act = run['results'].y[:9]
run['results'].y = y_act[:3]
fig = solver.plot3d(run)
run['results'].y = y_act[3:6]
fig = solver.plot3d(run, fig=fig)
run['results'].y = y_act[6:9]
fig = solver.plot3d(run, fig=fig)
plt.show(True)
def testExample(self):
sys = RestrictedCircular3Body()
slv = SystemSolver(sys)
lce, lce_run = slv.get_lce()
slv.calc_lce
# print(lce)
init = [1, -1, 1, -1]
run = slv.run([0, 5], init)
run['results'].y = run['results'].y[:2, :]
slv.plot2d(run)
class TestModel(unittest.TestCase):
def setUp(self):
# test data
self.circle_solver = SystemSolver(CircleSystem())
self.lorenz_solver = SystemSolver(LorenzSystem())
def test3dGraph(self):
res = self.lorenz_solver.run([0, 10], [1, 1, 1])
fig = self.lorenz_solver.plotnd(res)
self.lorenz_solver.plot3d(res)
self.assertEqual(4, len(fig.get_axes()))
def test2dGraph(self):
res = self.circle_solver.run([0, 10], [1, 1])
fig = self.circle_solver.plotnd(res)
self.circle_solver.plot2d(res)
self.assertEqual(3, len(fig.get_axes()))
def testMultiGraph(self):
run1 = self.lorenz_solver.run([0, 20], [1, 1, 1])
run2 = self.lorenz_solver.run([0, 20], [1, 1, 1])
run3 = self.lorenz_solver.run([0, 20], [1, 1, 1 + 10**-9])
fig = self.lorenz_solver.plot3d(run1)
fig = self.lorenz_solver.plot3d(run2, fig)
self.lorenz_solver.plot3d(run3, fig)
# You should see that orange (2nd graph) covers blue (1st graph) while
# adding a billionth to green (3rd graph) causes it to diverge.
def testMulti2dGraph(self):
run1 = self.circle_solver.run([0, 20], [0, 2])
run2 = self.circle_solver.run([0, 20], [0, 1])
fig = self.circle_solver.plot2d(run1)
self.circle_solver.plot2d(run2, fig)
def testMultiNdGraph(self):
run1 = self.lorenz_solver.run([0, 20], [1, 1, 1])
run2 = self.lorenz_solver.run([0, 20], [1, 1, 1.001])
fig = self.lorenz_solver.plotnd(run1)
self.lorenz_solver.plotnd(run2, fig)
self.k[i] *= self.e[i-1] * self.k[i-1]
self.levels = levels
def fun(self, t, v):
r = self.r
k = self.k
e = self.e
trans = r * v*v/k
y = np.zeros_like(v)
y[0] += r * v[0]
y[1:] += e*trans[:-1]
y -= trans
return y
if __name__ == '__main__':
biome1 = {'rate': 1, 'capacity': 1000, 'efficiency': [.1, .12], 'levels':3}
biome2 = {'rate': .1, 'capacity': 100, 'efficiency': [.1, .9], 'levels':3}
biome3 = {'rate': .01, 'capacity': 5000, 'efficiency': [.1, .2], 'levels':3}
problem = EcoSystem([biome1, biome2, biome3])
solver = SystemSolver(problem)
y0 = np.zeros(problem.dim)
y0[0] = 1
y0[3] = 1
y0[4] = 1000
y0[6] = 1
solver.run([0, 200], y0)
solver.plotnd()
plt.show(True)
import os
import numpy as np
import matplotlib.pyplot as plt
from echonn.sys import LorenzSystem, SystemSolver
if __name__ == "__main__":
lorenz_3d_plot = os.path.join('..', 'images', 'lorenz_3d_plot.png')
lorenz_nd_plot = os.path.join('..', 'images', 'lorenz_nd_plot.png')
slv = SystemSolver(LorenzSystem())
res = slv.run([0, 50], [10, 20, 30])
slv.plotnd(res, dims=['x', 'y', 'z'])
plt.savefig(lorenz_nd_plot)
slv.plot3d(res)
plt.savefig(lorenz_3d_plot)
plt.show(True)
y0 = np.zeros(2 * sys.body_dim * len(sys.body_masses), dtype=np.float64)
x1 = np.array([0.97000436, -0.24308753, 0])
x3p = np.array([-0.93240737, -0.86473146, 0])
y0[0:3] = x1
y0[3:6] = -x1
y0[6:9] = 0
y0[9:12] = -x3p / 2
y0[12:15] = -x3p / 2
y0[15:18] = x3p
tspan = [0, 100]
# print(sys.fun(np.zeros_like(y0), y0).reshape(6, -1))
run = solver.run(tspan, y0)
# solver.plotnd(run)
# print(run['results'].y[:, -1].reshape(6, -1))
y_act = run['results'].y[:9]
run['results'].y = y_act[:3]
fig = solver.plot3d(run)
run['results'].y = y_act[3:6]
fig = solver.plot3d(run, fig=fig)
run['results'].y = y_act[6:9]
fig = solver.plot3d(run, fig=fig)
plt.title('Long Run of 3 Body Figure 8')
plt.tight_layout()
plt.savefig(os.path.join(dir_pre, 'fig8_long.png'))
tspan = [0, 10]
# print(sys.fun(np.zeros_like(y0), y0).reshape(6, -1))
class TestDoubPenAnim(unittest.TestCase):
def setUp(self):
# test data
self.pendulum_system = DoublePendulumSystem()
self.pendulum_solver = SystemSolver(self.pendulum_system)
# run render test
self.render_til = 10
self.render_long = False # slows test a lot
self.render = False # slows test
def testWeirdLenSys(self):
sys = DoublePendulumSystem(1, 2, .5, 1.5)
slv = SystemSolver(sys)
tf = 3
run1 = slv.run([0, tf], [pi, 1, 1, 0])
run2 = slv.run([0, tf], [pi, 1.001, 1, 0])
if self.render:
self.runRender([run1, run2], 'odd_pend')
def testLongRun(self):
tf = 50
run1 = self.pendulum_solver.run([0, tf], [pi, 1, 1, 0])
run2 = self.pendulum_solver.run([0, tf], [pi, 1.001, 1, 0])
if self.render_long:
self.runRender([run1, run2], 'long_run', mult=.75)
def testRender(self):
run1 = self.pendulum_solver.run([0, self.render_til], [.2, 1, 1, 0])
run2 = self.pendulum_solver.run([0, self.render_til], [.2, 1.1, 1, 0])
if self.render:
self.runRender(run1, 'run1')
self.runRender([run1], 'run1_mult', mult=2)
self.runRender([run1, run2], 'run2', mult=.5)
def runRender(self, runs, fname, mult=1):
animator = DoublePendulumAnimator(runs, speed=mult)
animator.render()
fname = os.path.join('src', 'test', 'test_data', fname)
# http://thesmithfam.org/blog/2012/10/25/temporarily-suppress-console-output-in-python/
@contextmanager
def suppress_stdout():
with open(os.devnull, "w") as devnull:
old_stdout = sys.stdout
old_stderr = sys.stderr
sys.stdout = devnull
sys.stderr = devnull
try:
yield
finally:
sys.stdout = old_stdout
sys.stderr = old_stderr
try:
os.remove(fname + '.gif')
except:
pass # we don't care
self.assertFalse(os.path.isfile(fname + '.gif'))
with suppress_stdout():
animator.save(fname)
self.assertTrue(os.path.isfile(fname + '.gif'))
Arctic = 2000
Arid = 7000
def __init__(self, flight_rate=50, fire_rate=80, biome=None):
super().__init__(dim=2)
if biome is None:
biome = XwaModel.Grass
self.flight_rate = flight_rate
self.fire_rate = fire_rate
self.biome_carry_capacity = biome
self.biome_growth_rate = 10
def fun(self, t, v):
w, r = v
x = r * w
a = self.flight_rate + self.fire_rate
dw = (x - w * a) / 10**4
dr = -x + self.biome_growth_rate * r - self.biome_growth_rate * r**2 / self.biome_carry_capacity
if t < 100:
dw = 0
return dw, dr
if __name__ == '__main__':
problem = XwaModel()
solver = SystemSolver(problem)
y0 = np.array([8, 5000])
solver.run([0, 250], y0)
solver.plotnd()
plt.show(True)