def lorenz(cfg):
# Lorenz paramters and initial conditions
sigma, beta, rho = cfg.lorenz.sigma, cfg.lorenz.beta, cfg.lorenz.rho
u0, v0, w0 = cfg.lorenz.ex.u0, cfg.lorenz.ex.v0, cfg.lorenz.ex.w0
# Maximum time point and total number of time points
tmax, n = cfg.lorenz.tmax, cfg.lorenz.n
# Integrate the Lorenz equations on the time grid t
t = np.linspace(0, tmax, n)
f = odeint(sim_lorenz, (u0, v0, w0), t, args=(sigma, beta, rho))
x, y, z = f.T
num_traj = cfg.lorenz.num_traj
if cfg.collect_data:
data_X = np.zeros((1, 3))
data_Seq = []
new_init = np.random.uniform(low=[-25, -25, -25],
high=[25, 25, 25],
size=(num_traj, 3))
for row in new_init:
u, v, w = row # row[0], row[1], row[2]
f = odeint(sim_lorenz, (u, v, w), t, args=(sigma, beta, rho))
x, y, z = f.T
l = DotMap()
l.states = f
# Add parameters the way that the generation object is
# TODO take generic parameters rather than only PD Target
l.P = cfg.lorenz.beta
l.D = cfg.lorenz.rho
l.target = cfg.lorenz.sigma
data_Seq.append(l)
if cfg.plot: plot_lorenz(data_Seq, cfg, predictions=None)
if cfg.save_data:
log.info("Saving new default data")
torch.save((data_Seq),
hydra.utils.get_original_cwd() +
'/trajectories/lorenz/' + 'raw' + cfg.data_dir)
log.info(f"Saved trajectories to {cfg.data_dir}")
else:
data_Seq = torch.load(hydra.utils.get_original_cwd() +
'/trajectories/lorenz/' + 'raw' + cfg.data_dir)
# Analysis
from dynamics_model import DynamicsModel
from reacher_pd import create_dataset_step, create_dataset_traj
from plot import plot_loss
prob = cfg.model.prob
traj = cfg.model.traj
ens = cfg.model.ensemble
if traj:
dataset = create_dataset_traj(data_Seq, threshold=0.95)
else:
dataset = create_dataset_step(data_Seq)
if cfg.train_models:
model = DynamicsModel(cfg)
train_logs, test_logs = model.train(dataset, cfg)
plot_loss(train_logs,
test_logs,
cfg,
save_loc=cfg.env.name + '-' + cfg.model.str,
show=True)
if cfg.save_models:
log.info("Saving new default models")
torch.save(
model,
hydra.utils.get_original_cwd() + '/models/lorenz/' +
cfg.model.str + '.dat')
models = {}
for model_type in cfg.models_to_eval:
models[model_type] = torch.load(hydra.utils.get_original_cwd() +
'/models/lorenz/' + model_type +
".dat")
mse_evald = []
for i in range(cfg.num_eval):
traj_idx = np.random.randint(num_traj)
traj = data_Seq[traj_idx]
MSEs, predictions = test_models([traj], models)
MSE_avg = {key: np.average(MSEs[key], axis=0) for key in MSEs}
mse = {key: MSEs[key].squeeze() for key in MSEs}
mse_sub = {key: mse[key][mse[key] < 10**5] for key in mse}
pred = {key: predictions[key] for key in predictions}
mse_evald.append(mse)
#
# plot_states(traj.states, pred, save_loc="Predictions; traj-" + str(traj_idx), idx_plot=[0,1,2], show=False)
# plot_mse(mse_sub, save_loc="Error; traj-" + str(traj_idx), show=False)
plot_lorenz([traj], cfg, predictions=pred)
plot_mse_err(mse_evald, save_loc="Err Bar MSE of Predictions", show=True)
