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 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 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)
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)
score = []
higher_score = []
for rmse_res in results['best model rmse']:
ds_test, ys_test, total_rmse = rmse_res[3]
for sub in range(ds_test.shape[0]):
err = rmse(ds_test[:sub + 1], ys_test[:sub + 1])
if err > .05:
score.append(sub)
break
slvr = SystemSolver(run['system'])
runt = deepcopy(run)
runt['results'].t = ts_data.t
runt['results'].y = ts_data.y.T
fig = slvr.plotnd(runt, dims=['θ1', 'θ2', 'ω1', 'ω2'], overlay=False)
plt.savefig(os.path.join(dir_pre, 'full_differential.png'))
# plt.show(True)
runt['results'].t = ts_data.test_t
runt['results'].y = ds_test.T
slvr.plotnd(runt, dims=['θ1', 'θ2', 'ω1', 'ω2'], overlay=False)
plt.savefig(os.path.join(dir_pre, 'test_data.png'))
sorter = np.flip(np.argsort(score))
how_many = 5
for rank, i in enumerate(sorter[:how_many]):
ds_test, ys_test, total_rmse = results['best model rmse'][i][3]
print(i, total_rmse, results['params'][i])
rmse_over_t = [
rmse(ds_test[:sub + 1], ys_test[:sub + 1])
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)