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)
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
for sub in range(ds_test.shape[0]):
err = rmse(ds_test[:sub + 1], ys_test[:sub + 1])
if err > 1:
higher_score.append(sub)
break
runt = deepcopy(run)
runt['results'].t = ts_data.t
runt['results'].y = ts_data.y.T
fig = slvr.plot3d(runt)
plt.savefig(os.path.join(dir_pre, 'full_differential.png'))
sorter = np.flip(np.argsort(score))
how_many = 8
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])
for sub in range(ds_test.shape[0])
]
plt.figure()
plt.title('Test RMSE vs Lyapunov Time')
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[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))
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]