def test_smoother():
utils.set_rng_seed(1)
torch.set_default_dtype(torch.float64)
sigma = torch.tensor([10.])
rho = torch.tensor([28.])
beta = torch.tensor([8. / 3.])
C = torch.randn(2, 3)
dt = 0.04
sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint')
B = 1
T = 200
xs = [torch.randn(B, 1, 3)]
for t in range(T - 1):
xs.append(sys.step(torch.tensor([0.] * B), xs[-1]))
x = torch.cat(xs, dim=1).detach()
x.requires_grad = True
y = x.detach()
sys = BasicLagrangianSystem(qdim=3, dt=dt)
t = torch.stack([torch.arange(T),
torch.arange(T)]).to(torch.get_default_dtype())
x0 = torch.zeros_like(x)
x0.requires_grad = True
obscrit = GaussianObservationCriterion(torch.ones(3), t, y)
dyncrit = DELCriterion(t)
# Test GroupSOSCriterion
crit = GroupSOSCriterion([obscrit, dyncrit])
xsm, metrics = NLSsmoother(x0,
crit,
sys,
solver_kwargs={
'verbose': 2,
'tr_rho': 0.,
'max_nfev': 2
})
print('Passed.')
# Test BlockSparseGroupSOSCriterion
# crit = BlockSparseGroupSOSCriterion([obscrit, dyncrit])
# xsm, metrics = NLSsmoother(torch.zeros_like(x), crit, sys)
print('Passed.')
def test_learner():
utils.set_rng_seed(1)
torch.set_default_dtype(torch.float64)
sigma = torch.tensor([10.])
rho = torch.tensor([28.])
beta = torch.tensor([8. / 3.])
C = torch.randn(2, 3)
dt = 0.04
sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint')
B = 1
T = 200
xs = [torch.randn(B, 1, 3)]
for t in range(T - 1):
xs.append(sys.step(torch.tensor([0.] * B), xs[-1]))
x = torch.cat(xs, dim=1).detach()
t = torch.stack([torch.arange(T),
torch.arange(T)]).to(torch.get_default_dtype())
dyncrit = GaussianDynamicsCriterion(torch.ones(3), t)
params = list(sys.parameters())[:2]
vparams = parameters_to_vector(params)
true_vparams = vparams.clone()
vparams += torch.randn_like(vparams) * 0.01
vector_to_parameters(vparams, params)
opt_result = learner(sys, [dyncrit], [x], ['torch_minimize'], [params],
[{}],
opt_kwargs_list=[{
'nan_line_search': True,
'scheduler': 'ExponentialLR',
'scheduler_kwargs': {
'gamma': 0.99,
}
}])[0]
params = list(sys.parameters())[:2]
vparams = parameters_to_vector(params)
error = (vparams - true_vparams).norm().item()
assert np.allclose(error, 0., atol=1e-4), 'Error=%.3e' % error
print('Passed.')
def train_horizon_model(H, datadir):
# extract hyperparameters
hp = hyperparameters[H]
lr = hp['lr']
n_epochs = hp['n_epochs']
save_file = hp['save_file']
logdir = hp['logdir']
# load data
utils.set_rng_seed(1)
y_mean, y_std, u_mean, u_std = load_statistics(datadir)
train_data, train_trgt = load_helidata(datadir, 'train')
test_data, test_trgt = load_helidata(datadir, 'test')
valid_data, valid_trgt = load_helidata(datadir, 'valid')
# Define net
neural_net_kwargs = dict(input_size=10 * H,
hidden_sizes=[32] * 8,
output_size=6,
activation='tanh',
gain=1.0,
ln=False)
net = LagModel(neural_net_kwargs, H)
logger.setup(logdir, action='d')
# Train
system = train_net(net,
train_data,
train_trgt,
test_data,
test_trgt,
valid_data,
valid_trgt,
y_std,
lr,
logdir=logdir,
H=H,
n_epochs=n_epochs)
torch.save(system.state_dict(), save_file)
def main():
utils.set_rng_seed(1)
# load the data
for dset in ['0p01', '0p05', '0p10', '0p20', '0p30', '0p40', '0p50', '1p0']:
for damped in ['', 'damped_']:
data = torch.load('./datasets/%sdubpen_qddot.td'%damped)
data_ = torch.load('./datasets/%sdubpen_%s.td'%(damped,dset))
dt = 0.05
t_, q_, y = data_[:]
t, q, dq, ddq = data[:]
B, T, qdim = q.shape
with torch.no_grad():
std = 0.05 * 10
smoothed_q, smoothed_dq, smoothed_ddq = kalman_smooth_and_diff(y,dt, em_Q=False) # new!
if plot:
for b in range(B):
for n in range(2):
plt.subplot(3,2,1+n)
plt.plot(smoothed_q[b,:,n])
plt.plot(q_[b,:,n], '--')
plt.plot(y[b,:,n], alpha=0.5)
plt.subplot(3,2,3+n)
plt.plot(smoothed_dq[b,:,n])
plt.plot(dq[b,:,n], '--')
plt.subplot(3,2,5+n)
plt.plot(smoothed_ddq[b,:,n])
plt.plot(ddq[b,:,n], '--')
plt.show()
dataset = TensorDataset(t_, smoothed_q, smoothed_dq, smoothed_ddq)
torch.save(dataset, './datasets/%sdubpen_%s_smoothed.td'%(damped, dset))
def train(datadir, logdir, lr, n_epochs, H, save_file='lstm.th'):
utils.set_rng_seed(1)
y_mean, y_std, u_mean, u_std = load_statistics(datadir)
train_data, train_trgt = load_helidata(datadir, 'train')
test_data, test_trgt = load_helidata(datadir, 'test')
valid_data, valid_trgt = load_helidata(datadir, 'valid')
net = LSTMModel(**DEFAULT_LSTM_KWARGS)
logger.setup(logdir, action='d')
# Train
system = train_net(net,
train_data,
train_trgt,
test_data,
test_trgt,
valid_data,
valid_trgt,
y_std,
lr,
logdir=logdir,
H=H,
n_epochs=n_epochs)
torch.save(system.state_dict(), save_file)
def main(logdir, datafile, valdatafile, subset, obsmodel, practice, wfac, xdim, hotstartdir, seed,
parallel):
utils.set_rng_seed(seed)
torch.set_default_dtype(torch.float64)
logger.setup(logdir)
# load train data
traindata = np.load(datafile)
u = torch.tensor(traindata['u'])
y = torch.tensor(traindata['y'])
valdata = np.load(valdatafile)
val_u = torch.tensor(valdata['u'])
val_y = torch.tensor(valdata['y'])
_, y_std, _, _ = load_statistics('./datasets/split_normalized')
if practice == 1:
inds = np.random.choice(466, 10, replace=False)
u = u[inds]
y = y[inds]
B, T, ydim = y.shape
_, _, udim = u.shape
t = torch.stack([torch.arange(T)] * B).to(torch.get_default_dtype())
# specify system
sid = loadmat(opj(hotstartdir, 'SID_%dD.mat' % xdim))
A = torch.tensor(sid['A']).to(torch.get_default_dtype())
Bsys = torch.tensor(sid['B']).to(torch.get_default_dtype())
C = torch.tensor(sid['C']).to(torch.get_default_dtype())
D = torch.tensor(sid['D']).to(torch.get_default_dtype())
NN = deepcopy(DEFAULT_NN)
NN['gain'] = 0.1
sys = DiscreteLinear(xdim, udim, ydim, A=A, B=Bsys, C=C, D=D, obsModel=NN)
# specify smoothing criteria
smoothing_criteria = []
vfac = 1.0
for b in range(B):
obscrit = GaussianObservationCriterion(vfac * torch.ones(ydim), t[b:b + 1], y[b:b + 1],
u=u[b:b + 1])
dyncrit = GaussianDynamicsCriterion(wfac * torch.ones(xdim), t[b:b + 1], u=u[b:b + 1])
smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit]))
smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.1}
# specify learning criteria
learning_criteria = [
GaussianObservationCriterion(torch.ones(ydim), t, y, u=u),
GaussianDynamicsCriterion(torch.ones(xdim), t, u=u)
]
learning_params = [list(sys._obs.parameters()), list(sys._dyn.parameters())]
learning_opts = ['torch_minimize', 'torch_minimize']
learner_opt_kwargs = [{
'method': 'Adam',
'lr': 5e-4,
'nepochs': 500,
'tr_rho': 0.5
}, {
'method': 'Adam',
'lr': 1e-3,
'nepochs': 500,
'tr_rho': 0.5
}]
# save params
run_params = dict(seed=seed, subset=subset, xdim=xdim, vfac=vfac, wfac=wfac,
learning_opts=learning_opts, learner_opt_kwargs=learner_opt_kwargs,
smooth_solver_kwargs=smooth_solver_kwargs, practice=practice,
obsmodel=obsmodel)
with open(opj(logdir, 'run_params.json'), 'w') as f:
json.dump(run_params, f)
# instantiate CEEM
class Tracker:
def __init__(self):
return
tracker = Tracker()
tracker.best_val_rmse = np.inf
def ecb(epoch):
torch.save(sys.state_dict(),
os.path.join(logger.get_dir(), 'ckpts', 'model_{}.th'.format(epoch)))
y_pred = gen_ypred_model(sys, val_u, val_y)
rms = compute_rms(val_y[:, 25:], y_pred[:, 25:], y_std)
val_rmse = float(rms.mean())
logger.logkv('test/val_rmse', val_rmse)
if val_rmse < tracker.best_val_rmse:
tracker.best_val_rmse = val_rmse
torch.save(sys.state_dict(), os.path.join(logger.get_dir(), 'ckpts', 'best_model.th'))
return
epoch_callbacks = [ecb]
def tcb(epoch):
# TODO
return False
termination_callback = tcb
ceem = CEEM(smoothing_criteria, learning_criteria, learning_params, learning_opts,
epoch_callbacks, termination_callback, parallel=parallel)
# run CEEM
x0 = 0.01 * torch.randn(B, T, xdim)
ceem.train(xs=x0, sys=sys, nepochs=5000, smooth_solver_kwargs=smooth_solver_kwargs,
learner_opt_kwargs=learner_opt_kwargs, subset=subset)
def train(seed, logdir, sys_seed, k, b):
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
ystd = 0.01
# ystd = 0.
torch.set_default_dtype(torch.float64)
logger.setup(logdir, action='d')
N = 128
n = 3 * k
B = b
true_system = default_lorenz_system(k, obsdif=2)
utils.set_rng_seed(sys_seed)
xdim = true_system.xdim
ydim = true_system.ydim
dt = true_system._dt
x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k]).unsqueeze(0)
# simulate true_dynamics over IC distribution
x_test = x0mean.repeat(1024, 1, 1)
x_test += 5.0 * torch.randn_like(x_test)
x_test = x_test.detach()
t_test = torch.zeros(1024, 1)
tgt_test = true_system.step_derivs(t_test, x_test).detach()
## simulate the true system
xs = [x0mean.repeat(B, 1, 1)]
xs[0] += 2.5 * torch.randn_like(xs[0])
with torch.no_grad():
for t in range(N - 1):
xs.append(true_system.step(torch.tensor([0.] * B), xs[-1]))
xs = torch.cat(xs, dim=1)
fig = plt.figure()
for b in range(B):
ax = fig.add_subplot(int(np.ceil(B / 2.)), 2, b + 1, projection='3d')
for k_ in range(k):
plot3d(plt.gca(),
xs[b, :, k_],
xs[b, :, k + k_],
xs[b, :, 2 * k + k_],
linestyle='--',
alpha=0.5)
plt.savefig(os.path.join(logger.get_dir(), 'figs/traj_%d.png' % b),
dpi=300)
# plt.show()
plt.close()
t = torch.tensor(range(N)).unsqueeze(0).expand(B, -1).to(
torch.get_default_dtype())
y = true_system.observe(t, xs).detach()
# seed for real now
utils.set_rng_seed(seed)
y += ystd * torch.randn_like(y)
# prep system
system = deepcopy(true_system)
true_params = parameters_to_vector(true_system.parameters())
utils.set_rng_seed(seed)
params = true_params * (
(torch.rand_like(true_params) - 0.5) / 5. + 1.) # within 10%
vector_to_parameters(params, system.parameters())
params = list(system.parameters())
# specify smoothing criteria
smoothing_criteria = []
for b in range(B):
obscrit = GaussianObservationCriterion(1.0 * torch.ones(ydim),
t[b:b + 1], y[b:b + 1])
dyncrit = GaussianDynamicsCriterion(1e0 * torch.ones(xdim), t[b:b + 1])
smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit]))
smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.01}
# specify learning criteria
learning_criteria = [GaussianDynamicsCriterion(1e0 * torch.ones(xdim), t)]
learning_params = [params]
# learning_opts = ['scipy_minimize']
# learner_opt_kwargs = {'method': 'Nelder-Mead', 'tr_rho': 0.1,
# 'options':{'adaptive':True}}
# learner_opt_kwargs = {'method': 'BFGS', 'tr_rho': 0.1,
# 'options':{'disp':True}}
learning_opts = ['torch_minimize']
# learner_opt_kwargs = {
# 'method': 'Adam',
# 'lr': 5e-4,
# 'tr_rho': 0.1,
# 'nepochs': 200,
# 'max_grad_norm': 10.0
# }
learner_opt_kwargs = {'method': 'LBFGS'}
# instantiate CEEM
def ecb(epoch):
params = list(system.parameters())
vparams = parameters_to_vector(params)
error = (vparams - true_params).norm().item()
logger.logkv('test/log10_paramerror', np.log10(error))
return
epoch_callbacks = [ecb]
class Last10Errors:
def __init__(self):
return
last_10_errors = Last10Errors
last_10_errors._arr = []
def tcb(epoch):
with torch.no_grad():
tgt_test_pr = system.step_derivs(t_test, x_test)
error = float(torch.nn.functional.mse_loss(tgt_test_pr, tgt_test))
logger.logkv('test/log10_error', np.log10(error))
last_10_errors._arr.append(np.log10(error))
if len(last_10_errors._arr) > 100:
last_10_errors._arr = last_10_errors._arr[-100:]
l10err = torch.tensor(last_10_errors._arr)
convcrit = float((l10err.min() - l10err.max()).abs())
logger.logkv('test/log10_convcrit', np.log10(convcrit))
if convcrit < 1e-3:
return True
return False
termination_callback = tcb
ecb(-1)
tcb(-1)
logger.dumpkvs()
ceem = CEEM(smoothing_criteria,
learning_criteria,
learning_params,
learning_opts,
epoch_callbacks,
termination_callback,
parallel=min(4, B))
# run CEEM
# x0 = torch.zeros_like(xs)
x0 = xs + torch.randn_like(xs)
ceem.train(xs=x0,
sys=system,
nepochs=5000,
smooth_solver_kwargs=smooth_solver_kwargs,
learner_opt_kwargs=learner_opt_kwargs)
return
def main():
utils.set_rng_seed(2)
torch.set_default_dtype(torch.float64)
dt = 0.05
sys = LagrangianDoublePendulum(dt, 1., 1., 1., 1., 10.)
q1 = torch.rand(16, 1, 2) * np.pi - np.pi / 2
q2 = q1.clone()
qs = [q1, q2]
for t in tqdm(range(200)):
qt = qs[-2].detach()
qtp1 = qs[-1].detach()
nq = sys.variational_step(qt, qtp1, oneatatime=True)
qs.append(nq)
x = torch.cat(qs, dim=1).detach()
B, T, _ = x.shape
t = torch.arange(T).unsqueeze(0).repeat(B, 1).float()
y_p01 = (x + 0.01 * torch.randn_like(x)).detach()
y_p05 = (x + 0.05 * torch.randn_like(x)).detach()
y_p10 = (x + 0.1 * torch.randn_like(x)).detach()
y_p20 = (x + 0.2 * torch.randn_like(x)).detach()
y_p30 = (x + 0.3 * torch.randn_like(x)).detach()
y_p40 = (x + 0.4 * torch.randn_like(x)).detach()
y_p50 = (x + 0.5 * torch.randn_like(x)).detach()
y_1p0 = (x + 1.0 * torch.randn_like(x)).detach()
dataset = TensorDataset(t, x, y_p01)
torch.save(dataset, './datasets/dubpen_0p01.td')
dataset = TensorDataset(t, x, y_p05)
torch.save(dataset, './datasets/dubpen_0p05.td')
dataset = TensorDataset(t, x, y_p10)
torch.save(dataset, './datasets/dubpen_0p10.td')
dataset = TensorDataset(t, x, y_p20)
torch.save(dataset, './datasets/dubpen_0p20.td')
dataset = TensorDataset(t, x, y_p30)
torch.save(dataset, './datasets/dubpen_0p30.td')
dataset = TensorDataset(t, x, y_p40)
torch.save(dataset, './datasets/dubpen_0p40.td')
dataset = TensorDataset(t, x, y_p50)
torch.save(dataset, './datasets/dubpen_0p50.td')
dataset = TensorDataset(t, x, y_1p0)
torch.save(dataset, './datasets/dubpen_1p0.td')
x_ = (x[:, 1:] + x[:, :-1]) * 0.5
dx = (x[:, 1:] - x[:, :-1]) / dt
x = x_
ddx = sys.compute_qddot(x, dx)
B, T, _ = x.shape
t = torch.arange(T).unsqueeze(0).repeat(B, 1).float()
dataset = TensorDataset(t, x, dx, ddx)
torch.save(dataset, './datasets/dubpen_qddot.td')
if plot:
for b in range(16):
plt.subplot(8, 2, b + 1)
plt.plot(x[b])
plt.show()
def test_particleem():
utils.set_rng_seed(1)
torch.set_default_dtype(torch.float64)
# setup system
sigma = torch.tensor([10.])
rho = torch.tensor([28.])
beta = torch.tensor([8. / 3.])
C = torch.randn(2, 3)
dt = 0.04
sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint')
# simulate
B = 2
T = 20
xs = [torch.randn(B, 1, 3)]
for t in range(T - 1):
xs.append(sys.step(torch.tensor([0.] * B), xs[-1]))
xtr = torch.cat(xs, dim=1).detach()
y = sys.observe(0., xtr).detach()
# y += torch.rand_like(y) * 0.01
t = torch.stack([torch.arange(T)] * B).to(torch.get_default_dtype())
#
params = list(sys.parameters())
vparams = parameters_to_vector(params)
true_vparams = vparams.clone()
vparams *= 1. + (torch.rand_like(vparams) - 0.5) * 2 * 0.025
vector_to_parameters(vparams, params)
Px0 = torch.eye(3)
Q = 0.1 * torch.eye(3)
R = 0.1 * torch.eye(2)
Np = 300
fapf = faPF(Np, sys, Q, R, Px0)
def callback(epoch):
pass
trainer = SAEMTrainer(
fapf,
y,
gamma_sched=lambda x: 0.8,
xlen_cutoff=3,
max_k=10,
)
trainer.train(params, callbacks=[callback])
assert True
if True:
print('Passed.')
else:
print('Failed.')
def test_smoother():
utils.set_rng_seed(1)
torch.set_default_dtype(torch.float64)
sigma = torch.tensor([10.])
rho = torch.tensor([28.])
beta = torch.tensor([8. / 3.])
C = torch.randn(2, 3)
dt = 0.04
sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint')
B = 1
T = 200
xs = [torch.randn(B, 1, 3)]
for t in range(T - 1):
xs.append(sys.step(torch.tensor([0.] * B), xs[-1]))
x = torch.cat(xs, dim=1).detach()
x.requires_grad = True
y = sys.observe(0., x).detach()
# y += torch.rand_like(y) * 0.01
t = torch.stack([torch.arange(T),
torch.arange(T)]).to(torch.get_default_dtype())
x0 = torch.zeros_like(x)
obscrit = GaussianObservationCriterion(torch.ones(2), t, y)
dyncrit = GaussianDynamicsCriterion(torch.ones(3), t)
# Test GroupSOSCriterion
crit = GroupSOSCriterion([obscrit, dyncrit])
xsm, metrics = NLSsmoother(x0,
crit,
sys,
solver_kwargs={
'verbose': 2,
'tr_rho': 0.
})
err = float((xsm - x).norm())
assert err < 1e-8, 'Smoothing Error: %.3e' % err
print('Passed.')
# Test BlockSparseGroupSOSCriterion
crit = BlockSparseGroupSOSCriterion([obscrit, dyncrit])
xsm, metrics = NLSsmoother(torch.zeros_like(x), crit, sys)
err = float((xsm - x).norm())
assert err < 1e-8, 'Smoothing Error: %.3e' % err
print('Passed.')
def train(seed, logdir, ystd=0.1, wstd=0.01, sys_seed=4):
print('\n\n\n##### SEED %d #####\n\n' % seed)
torch.set_default_dtype(torch.float64)
logger.setup(logdir, action='d')
# Number of timesteps in the trajectory
T = 128
n = 3
# Batch size
B = 4
k = 1
utils.set_rng_seed(sys_seed)
sys = default_lorenz_attractor()
dt = sys._dt
utils.set_rng_seed(seed)
# simulate the system
x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k])
x0mean = x0mean.unsqueeze(0).repeat(B, 1, 1)
# Rollout with noise
Q = (wstd**2) * torch.eye(sys.xdim)
R = (ystd**2) * torch.eye(sys.ydim)
Px0 = 5.0 * torch.eye(sys.xdim)
Qpdf = MultivariateNormal(torch.zeros((B, 1, sys.xdim)),
Q.unsqueeze(0).unsqueeze(0))
Rpdf = MultivariateNormal(torch.zeros((B, 1, sys.ydim)),
R.unsqueeze(0).unsqueeze(0))
Px0pdf = MultivariateNormal(x0mean, Px0.unsqueeze(0).unsqueeze(0))
xs = [Px0pdf.sample()]
ys = [sys.observe(0, xs[0]) + Rpdf.sample()]
for t in range(T - 1):
tinp = torch.tensor([t] *
B).unsqueeze(1).to(dtype=torch.get_default_dtype())
xs.append(sys.step(tinp, xs[-1]) + Qpdf.sample())
ys.append(sys.observe(tinp, xs[-1]) + Rpdf.sample())
x = torch.cat(xs, dim=1)
ys = torch.cat(ys, dim=1)
m = ys.shape[-1]
fig = plt.figure()
for b in range(B):
ax = fig.add_subplot(int(np.ceil(B / 2.)), 2, b + 1, projection='3d')
for k_ in range(k):
plot3d(plt.gca(),
x[b, :, k_],
x[b, :, k + k_],
x[b, :, 2 * k + k_],
linestyle='--',
alpha=0.5)
plt.savefig(opj(logdir, 'traintrajs.png'), dpi=300)
true_system = deepcopy(sys)
Np = 100
params = [sys._sigma, sys._rho, sys._beta]
true_vparams = parameters_to_vector(params)
pert_vparams = true_vparams * (
(torch.rand_like(true_vparams) - 0.5) / 5 + 1.0)
vector_to_parameters(pert_vparams, params)
fapf = faPF(Np, sys, Q, R, Px0)
timer = {'start_time': timeit.default_timer()}
def callback(epoch):
logger.logkv('test/rho', float(sys._rho))
logger.logkv('test/sigma', float(sys._sigma))
logger.logkv('test/beta', float(sys._beta))
logger.logkv(
'test/rho_pcterr_log10',
float(
torch.log10(
(true_system._rho - sys._rho).abs() / true_system._rho)))
logger.logkv(
'test/sigma_pcterr_log10',
float(
torch.log10((true_system._sigma - sys._sigma).abs() /
true_system._sigma)))
logger.logkv(
'test/beta_pcterr_log10',
float(
torch.log10((true_system._beta - sys._beta).abs() /
true_system._beta)))
logger.logkv('time/epochtime',
timeit.default_timer() - timer['start_time'])
timer['start_time'] = timeit.default_timer()
return
callback(-1)
logger.dumpkvs()
trainer = SAEMTrainer(
fapf,
ys,
# gamma_sched=lambda x: HarmonicDecayScheduler(x, a=50.),
gamma_sched=lambda x: 0.8)
trainer.train(params, callbacks=[callback])
def train(seed, logdir, sys_seed, ystd, wstd):
torch.set_default_dtype(torch.float64)
logger.setup(logdir, action='d')
# Number of timesteps in the trajectory
T = 128
n = 3
# Batch size
B = 1
k = 1
utils.set_rng_seed(sys_seed)
true_system = default_lorenz_attractor()
dt = true_system._dt
utils.set_rng_seed(43)
# simulate the system
x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k]).unsqueeze(0)
# seed for real now
utils.set_rng_seed(seed)
# Rollout with noise
xs = [x0mean]
xs[0] += 5. * torch.randn_like(xs[0])
with torch.no_grad():
for t in range(T - 1):
xs.append(
true_system.step(torch.tensor([0.] * B), xs[-1]) + wstd * torch.randn_like(xs[-1]))
xs = torch.cat(xs, dim=1)
t = torch.tensor(range(T)).unsqueeze(0).to(torch.get_default_dtype())
y = true_system.observe(t, xs).detach()
y += ystd * torch.randn_like(y) # Observation noise
# prep system
system = deepcopy(true_system)
true_params = parameters_to_vector(true_system.parameters())
params = true_params * ((torch.rand_like(true_params) - 0.5) / 5. + 1.) # within 10%
vector_to_parameters(params, system.parameters())
params = list(system.parameters())
# specify smoothing criteria
B = 1
smoothing_criteria = []
for b in range(B):
obscrit = GaussianObservationCriterion(torch.ones(2), t[b:b + 1], y[b:b + 1])
dyncrit = GaussianDynamicsCriterion(wstd / ystd * torch.ones(3), t[b:b + 1])
smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit]))
smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.001}
# specify learning criteria
learning_criteria = [GaussianDynamicsCriterion(torch.ones(3), t)]
learning_params = [params]
learning_opts = ['scipy_minimize']
learner_opt_kwargs = {'method': 'Nelder-Mead', 'tr_rho': 0.01}
# instantiate CEEM
def ecb(epoch):
logger.logkv('test/rho', float(system._rho))
logger.logkv('test/sigma', float(system._sigma))
logger.logkv('test/beta', float(system._beta))
logger.logkv('test/rho_pcterr_log10',
float(torch.log10((true_system._rho - system._rho).abs() / true_system._rho)))
logger.logkv(
'test/sigma_pcterr_log10',
float(torch.log10((true_system._sigma - system._sigma).abs() / true_system._sigma)))
logger.logkv(
'test/beta_pcterr_log10',
float(torch.log10((true_system._beta - system._beta).abs() / true_system._beta)))
return
epoch_callbacks = [ecb]
class Last10Errors:
def __init__(self):
return
last_10_errors = Last10Errors
last_10_errors._arr = []
def tcb(epoch):
params = list(system.parameters())
vparams = parameters_to_vector(params)
error = (vparams - true_params).norm().item()
last_10_errors._arr.append(float(error))
logger.logkv('test/log10_error', np.log10(error))
if len(last_10_errors._arr) > 10:
last_10_errors._arr = last_10_errors._arr[-10:]
l10err = torch.tensor(last_10_errors._arr)
convcrit = float((l10err.min() - l10err.max()).abs())
logger.logkv('test/log10_convcrit', np.log10(convcrit))
if convcrit < 1e-4:
return True
return False
termination_callback = tcb
ceem = CEEM(smoothing_criteria, learning_criteria, learning_params, learning_opts,
epoch_callbacks, termination_callback)
# run CEEM
x0 = torch.zeros_like(xs)
ceem.train(xs=x0, sys=system, nepochs=500, smooth_solver_kwargs=smooth_solver_kwargs,
learner_opt_kwargs=learner_opt_kwargs)
return float(system._sigma), float(system._rho), float(system._beta)
def test_sys():
utils.set_rng_seed(1)
torch.set_default_dtype(torch.float64)
# test LorenzAttractor
sigma = torch.tensor([10.])
rho = torch.tensor([28.])
beta = torch.tensor([8. / 3.])
C = torch.randn(2, 3)
dt = 0.04
sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint')
B = 5
T = 20
xs = [torch.randn(B, 1, 3)]
for t in range(T - 1):
xs.append(sys.step(torch.tensor([0.] * B), xs[-1]))
x = torch.cat(xs, dim=1).detach()
x.requires_grad = True
y = sys.observe(0., x)
y += torch.rand_like(y) * 0.01
t = torch.stack([torch.arange(T),
torch.arange(T)]).to(torch.get_default_dtype())
check_sys(sys, t, x, y)
# test SpringMassDamper
n = 4
M = D = K = torch.tensor([[1., 2.], [2., 5.]])
dt = 0.1
method = 'midpoint'
sys = SpringMassDamper(M, D, K, dt, method=method)
B = 5
T = 20
xs = [torch.randn(B, 1, n)]
for t in range(T - 1):
xs.append(sys.step(torch.tensor([0.] * B), xs[-1]))
x = torch.cat(xs, dim=1).detach()
x.requires_grad = True
y = sys.observe(0., x).detach()
y += torch.rand_like(y) * 0.01
t = torch.stack([torch.arange(T),
torch.arange(T)]).to(torch.get_default_dtype())
check_sys(sys, t, x, y)
# test DiscreteLinear
xdim = 2
ydim = 3
udim = 2
sys = DiscreteLinear(xdim, udim, ydim)
x = torch.randn(B, T, xdim)
u = torch.randn(B, T, udim)
y = torch.randn(B, T, ydim)
t = torch.stack([torch.arange(T),
torch.arange(T)]).to(torch.get_default_dtype())
check_sys(sys, t, x, y, u=u)
def test_ceem():
utils.set_rng_seed(1)
torch.set_default_dtype(torch.float64)
# setup system
sigma = torch.tensor([10.])
rho = torch.tensor([28.])
beta = torch.tensor([8. / 3.])
C = torch.randn(2, 3)
dt = 0.04
sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint')
# simulate
B = 2
T = 20
xs = [torch.randn(B, 1, 3)]
for t in range(T - 1):
xs.append(sys.step(torch.tensor([0.] * B), xs[-1]))
xtr = torch.cat(xs, dim=1).detach()
y = sys.observe(0., xtr).detach()
# y += torch.rand_like(y) * 0.01
t = torch.stack([torch.arange(T)] * B).to(torch.get_default_dtype())
#
params = list(sys.parameters())
vparams = parameters_to_vector(params)
true_vparams = vparams.clone()
vparams *= 1. + (torch.rand_like(vparams) - 0.5) * 2 * 0.025
vector_to_parameters(vparams, params)
# specify smoothing criteria
smoothing_criteria = []
for b in range(B):
obscrit = GaussianObservationCriterion(torch.ones(2), t[b:b + 1],
y[b:b + 1])
dyncrit = GaussianDynamicsCriterion(torch.ones(3), t[b:b + 1])
smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit]))
smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.001}
# specify learning criteria
learning_criteria = [GaussianDynamicsCriterion(torch.ones(3), t)]
learning_params = [params]
learning_opts = ['scipy_minimize']
learner_opt_kwargs = {'method': 'Nelder-Mead', 'tr_rho': 0.01}
# learning_opts = ['torch_minimize']
# learner_opt_kwargs = {'method': 'Adam',
# 'tr_rho': 0.01}
# instantiate CEEM
def ecb(epoch):
return
epoch_callbacks = [ecb]
def tcb(epoch):
params = list(sys.parameters())
vparams = parameters_to_vector(params)
error = (vparams - true_vparams).norm().item()
logger.logkv('test/log10_error', np.log10(error))
return error < 5e-3
termination_callback = tcb
ceem = CEEM(smoothing_criteria,
learning_criteria,
learning_params,
learning_opts,
epoch_callbacks,
termination_callback,
parallel=2)
hp_scheduler = {'tr_rho': utils.LambdaScheduler(0.01, lambda x: x)}
# run CEEM
x0 = torch.zeros_like(xtr)
ceem.train(xs=x0,
sys=sys,
nepochs=150,
smooth_solver_kwargs=smooth_solver_kwargs,
learner_opt_kwargs=learner_opt_kwargs,
subset=1,
hp_schedulers=hp_scheduler)
assert tcb(0)
if tcb(0):
print('Passed.')
else:
print('Failed.')
def train(seed, logdir, ystd=0.1, wstd=0.01, sys_seed=4):
print('\n\n\n##### SEED %d #####\n\n'%seed)
torch.set_default_dtype(torch.float64)
logger.setup(logdir, action='d')
# Number of timesteps in the trajectory
T = 128
n = 3
# Batch size
B = 4
k = 1
utils.set_rng_seed(sys_seed)
sys = default_lorenz_attractor()
dt = sys._dt
utils.set_rng_seed(seed)
# simulate the system
x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k])
x0mean = x0mean.unsqueeze(0).repeat(B,1,1)
# Rollout with noise
Q = (wstd ** 2) * torch.eye(sys.xdim)
R = (ystd ** 2) * torch.eye(sys.ydim)
Px0 = 5.0 * torch.eye(sys.xdim)
Qpdf = MultivariateNormal(torch.zeros((B,1,sys.xdim)), Q.unsqueeze(0).unsqueeze(0))
Rpdf = MultivariateNormal(torch.zeros((B,1,sys.ydim)), R.unsqueeze(0).unsqueeze(0))
Px0pdf = MultivariateNormal(x0mean, Px0.unsqueeze(0).unsqueeze(0))
xs = [Px0pdf.sample()]
ys = [sys.observe(0, xs[0]) + Rpdf.sample()]
for t in range(T-1):
tinp = torch.tensor([t] * B).unsqueeze(1).to(dtype=torch.get_default_dtype())
xs.append(sys.step(tinp, xs[-1]) + Qpdf.sample())
ys.append(sys.observe(tinp, xs[-1]) + Rpdf.sample())
x = torch.cat(xs, dim=1)
y = torch.cat(ys, dim=1)
t = torch.tensor(range(T)).unsqueeze(0).to(torch.get_default_dtype()).repeat(B,1)
m = y.shape[-1]
fig = plt.figure()
for b in range(B):
ax = fig.add_subplot(int(np.ceil(B / 2.)), 2, b + 1, projection='3d')
for k_ in range(k):
plot3d(plt.gca(), x[b, :, k_], x[b, :, k + k_], x[b, :, 2 * k + k_], linestyle='--',
alpha=0.5)
plt.savefig(opj(logdir, 'traintrajs.png'), dpi=300)
# prep system
true_system = sys
system = deepcopy(true_system)
true_params = parameters_to_vector(true_system.parameters())
params = true_params * ((torch.rand_like(true_params) - 0.5) / 5. + 1.) # within 10%
vector_to_parameters(params, system.parameters())
params = list(system.parameters())
# specify smoothing criteria
smoothing_criteria = []
for b in range(B):
obscrit = GaussianObservationCriterion(torch.ones(2), t[b:b + 1], y[b:b + 1])
dyncrit = GaussianDynamicsCriterion(wstd / ystd * torch.ones(3), t[b:b + 1])
smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit]))
smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.001}
# specify learning criteria
learning_criteria = [GaussianDynamicsCriterion(torch.ones(3), t)]
learning_params = [params]
learning_opts = ['scipy_minimize']
learner_opt_kwargs = {'method': 'Nelder-Mead', 'tr_rho': 0.01}
# instantiate CEEM
timer = {'start_time':timeit.default_timer()}
def ecb(epoch):
logger.logkv('test/rho', float(system._rho))
logger.logkv('test/sigma', float(system._sigma))
logger.logkv('test/beta', float(system._beta))
logger.logkv('test/rho_pcterr_log10',
float(torch.log10((true_system._rho - system._rho).abs() / true_system._rho)))
logger.logkv(
'test/sigma_pcterr_log10',
float(torch.log10((true_system._sigma - system._sigma).abs() / true_system._sigma)))
logger.logkv(
'test/beta_pcterr_log10',
float(torch.log10((true_system._beta - system._beta).abs() / true_system._beta)))
logger.logkv('time/epochtime', timeit.default_timer() - timer['start_time'])
timer['start_time'] = timeit.default_timer()
return
epoch_callbacks = [ecb]
class Last10Errors:
def __init__(self):
return
last_10_errors = Last10Errors
last_10_errors._arr = []
def tcb(epoch):
params = list(system.parameters())
vparams = parameters_to_vector(params)
error = (vparams - true_params).norm().item()
last_10_errors._arr.append(float(error))
logger.logkv('test/log10_error', np.log10(error))
if len(last_10_errors._arr) > 10:
last_10_errors._arr = last_10_errors._arr[-10:]
l10err = torch.tensor(last_10_errors._arr)
convcrit = float((l10err.min() - l10err.max()).abs())
logger.logkv('test/log10_convcrit', np.log10(convcrit))
if convcrit < 1e-4:
return True
return False
termination_callback = tcb
ceem = CEEM(smoothing_criteria, learning_criteria, learning_params, learning_opts,
epoch_callbacks, termination_callback)
# run CEEM
x0 = torch.zeros_like(x)
ecb(-1)
logger.dumpkvs()
ceem.train(xs=x0, sys=system, nepochs=100, smooth_solver_kwargs=smooth_solver_kwargs,
learner_opt_kwargs=learner_opt_kwargs)
return float(system._sigma), float(system._rho), float(system._beta)
def train(seed, logdir, sys_seed, k, b):
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
ystd = 0.01
# ystd = 0.
torch.set_default_dtype(torch.float64)
logger.setup(logdir, action='d')
N = 128
n = 3 * k
B = b
true_system = default_lorenz_system(k, obsdif=2)
utils.set_rng_seed(sys_seed)
xdim = true_system.xdim
ydim = true_system.ydim
dt = true_system._dt
x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k]).unsqueeze(0)
# simulate true_dynamics over IC distribution
x_test = x0mean.repeat(1024, 1, 1)
x_test += 5.0 * torch.randn_like(x_test)
x_test = x_test.detach()
t_test = torch.zeros(1024, 1)
tgt_test = true_system.step_derivs(t_test, x_test).detach()
Q = (0.01**2) * torch.eye(true_system.xdim)
R = (ystd**2) * torch.eye(true_system.ydim)
Px0 = 2.5**2 * torch.eye(true_system.xdim)
## simulate the true system
xs = [x0mean.repeat(B, 1, 1)]
xs[0] += 2.5 * torch.randn_like(xs[0])
with torch.no_grad():
for t in range(N - 1):
xs.append(true_system.step(torch.tensor([0.] * B), xs[-1]))
xs = torch.cat(xs, dim=1)
fig = plt.figure()
for b in range(B):
ax = fig.add_subplot(int(np.ceil(B / 2.)), 2, b + 1, projection='3d')
for k_ in range(k):
plot3d(plt.gca(),
xs[b, :, k_],
xs[b, :, k + k_],
xs[b, :, 2 * k + k_],
linestyle='--',
alpha=0.5)
plt.savefig(os.path.join(logger.get_dir(), 'figs/traj_%d.png' % b),
dpi=300)
# plt.show()
plt.close()
t = torch.tensor(range(N)).unsqueeze(0).expand(B, -1).to(
torch.get_default_dtype())
y = true_system.observe(t, xs).detach()
# seed for real now
utils.set_rng_seed(seed)
y += ystd * torch.randn_like(y)
# prep system
system = deepcopy(true_system)
true_params = parameters_to_vector(true_system.parameters())
utils.set_rng_seed(seed)
params = true_params * (
(torch.rand_like(true_params) - 0.5) / 5. + 1.) # within 10%
vector_to_parameters(params, system.parameters())
params = list(system.parameters())
Np = 100
fapf = faPF(Np, system, Q, R, Px0)
timer = {'start_time': timeit.default_timer()}
def ecb(epoch):
logger.logkv('time/epoch', epoch)
params = list(system.parameters())
vparams = parameters_to_vector(params)
error = (vparams - true_params).norm().item()
logger.logkv('test/log10_paramerror', np.log10(error))
logger.logkv('time/epochtime',
timeit.default_timer() - timer['start_time'])
timer['start_time'] = timeit.default_timer()
with torch.no_grad():
tgt_test_pr = system.step_derivs(t_test, x_test)
error = float(torch.nn.functional.mse_loss(tgt_test_pr, tgt_test))
logger.logkv('test/log10_error', np.log10(error))
return
epoch_callbacks = [ecb]
ecb(-1)
logger.dumpkvs()
trainer = SAEMTrainer(
fapf,
y,
# gamma_sched=lambda x: HarmonicDecayScheduler(x, a=50.),
gamma_sched=lambda x: 0.2,
max_k=5000,
xlen_cutoff=15,
)
trainer.train(params, callbacks=epoch_callbacks)
return
