np.random.seed(args.seed)
num_stages = 3
num_layers = np.random.randint(1, 3)
num_units = np.random.randint(10, 25)
init_params = {'loc': 0., 'scale': 3.}
q_theta, opt_data, save_path, failed = M.epi(
mu,
arch_type='coupling',
num_stages=num_stages,
num_layers=num_layers,
num_units=num_units,
post_affine=False,
init_params=init_params,
K=10,
num_iters=1000,
N=500,
lr=1e-3,
c0=1e-3,
verbose=True,
stop_early=True,
log_rate=50,
save_movie_data=True,
)
print("EPI done.")
print("Saved to %s." % save_path)
if not failed:
print("Writing movie...")
M.epi_opt_movie(save_path)
print("done.")
# 3. Run EPI.
init_params = {'loc': 0., 'scale': 2.}
q_theta, opt_data, save_path, failed = model.epi(
mu,
arch_type='coupling',
num_stages=num_stages,
num_layers=2,
num_units=num_units,
post_affine=True,
batch_norm=True,
init_params=init_params,
K=15,
N=500,
num_iters=5000,
lr=1e-3,
c0=c0,
beta=4.,
nu=1.0,
random_seed=random_seed,
verbose=True,
stop_early=True,
log_rate=250,
save_movie_data=True,
)
if not failed:
print("Making movie.")
model.epi_opt_movie(save_path)
print("done.")
def test_epi():
mu = np.array([0.0, 0.1, 2 * np.pi, 0.1 * np.pi])
lb_a12 = 0.0
ub_a12 = 10.0
lb_a21 = -10.0
ub_a21 = 0.0
a11 = Parameter("a11", 1, 0.0)
a12 = Parameter("a12", 1, lb_a12, ub_a12)
a21 = Parameter("a21", 1, lb_a21, ub_a21)
a22 = Parameter("a22", 1, ub=0.0)
params = [a11, a12, a21, a22]
M = Model("lds", params)
M.set_eps(linear2D_freq)
q_theta, opt_data, save_path, _ = M.epi(
mu, num_iters=100, K=1, save_movie_data=True
)
g = q_theta.plot_dist()
M.epi_opt_movie(save_path)
params = [a11, a12, a21, a22]
M = Model("lds_2D", params)
M.set_eps(linear2D_freq)
q_theta, opt_data, save_path, _ = M.epi(
mu, num_iters=100, K=1, save_movie_data=True
)
q_theta = M.load_epi_dist(mu, k=1)
M.epi_opt_movie(save_path)
q_theta, opt_data, save_path, _ = M.epi(
mu, num_units=31, num_iters=100, K=1, save_movie_data=True
)
M.plot_epi_hpsearch(mu)
opt_data_filename = save_path + "opt_data.csv"
opt_data_cols = ["k", "iteration", "H", "converged"] + [
"R%d" % i for i in range(1, M.m + 1)
]
for x, y in zip(opt_data.columns, opt_data_cols):
assert x == y
# opt_data_df = pd.read_csv(opt_data_filename)
# opt_data_df['iteration'] = 2*opt_data_df['iteration']
# opt_data_df.to_csv(opt_data_filename)
# with raises(IOError):
# M.epi_opt_movie(save_path)
# os.remove(opt_data_filename)
# with raises(IOError):
# M.epi_opt_movie(save_path)
assert q_theta is not None
with raises(ValueError):
q_theta = M.load_epi_dist(mu, k=20)
with raises(TypeError):
q_theta = M.load_epi_dist(mu, k="foo")
with raises(ValueError):
q_theta = M.load_epi_dist(mu, k=-1)
M = Model("foo", params)
with raises(ValueError):
q_theta = M.load_epi_dist(mu, k=-1)
z = q_theta(1000)
log_q_z = q_theta.log_prob(z)
assert np.sum(z[:, 0] < 0.0) == 0
assert np.sum(z[:, 1] < lb_a12) == 0
assert np.sum(z[:, 1] > ub_a12) == 0
assert np.sum(z[:, 2] < lb_a21) == 0
assert np.sum(z[:, 2] > ub_a21) == 0
assert np.sum(z[:, 3] > 0.0) == 0
assert np.sum(1 - np.isfinite(z)) == 0
assert np.sum(1 - np.isfinite(log_q_z)) == 0
# Intentionally swap order in list to insure proper handling.
params = [a22, a21, a12, a11]
M = Model("lds2", params)
M.set_eps(linear2D_freq)
q_theta, opt_data, save_path, _ = M.epi(
mu, K=2, num_iters=100, stop_early=True, verbose=True
)
with raises(IOError):
M.epi_opt_movie(save_path)
z = q_theta(1000)
log_q_z = q_theta.log_prob(z)
assert np.sum(z[:, 0] < 0.0) == 0
assert np.sum(z[:, 1] < lb_a12) == 0
assert np.sum(z[:, 1] > ub_a12) == 0
assert np.sum(z[:, 2] < lb_a21) == 0
assert np.sum(z[:, 2] > ub_a21) == 0
assert np.sum(z[:, 3] > 0.0) == 0
assert np.sum(1 - np.isfinite(z)) == 0
assert np.sum(1 - np.isfinite(log_q_z)) == 0
for x, y in zip(opt_data.columns, opt_data_cols):
assert x == y
with raises(ValueError):
def bad_f(a11, a12, a21, a22):
return tf.expand_dims(a11 + a12 + a21 + a22, 0)
M.set_eps(bad_f)
params = [a22, a21, a12, a11]
M = Model("lds2", params)
nf = NormalizingFlow("autoregressive", 4, 1, 2, 10)
al_hps = AugLagHPs()
with raises(AttributeError):
save_path = M.get_save_path(mu, nf, al_hps, None)
save_path = M.get_save_path(mu, nf, al_hps, eps_name="foo")
return None
q_theta, opt_data, epi_path, failed = model.epi(
mu,
arch_type="coupling",
num_stages=2,
num_layers=num_layers,
num_units=25,
post_affine=True,
elemwise_fn="affine",
batch_norm=False,
bn_momentum=0.0,
K=6,
N=400,
num_iters=5000,
lr=1e-3,
c0=c0,
beta=beta,
nu=0.5,
random_seed=random_seed,
init_type=init_type,
init_params=init_params,
verbose=True,
stop_early=True,
log_rate=50,
save_movie_data=True,
)
if not failed:
print("Making movie.", flush=True)
model.epi_opt_movie(epi_path)
print("done.", flush=True)
def test_epi():
mu = np.array([0.0, 0.1, 2 * np.pi, 0.1 * np.pi])
lb_a12 = 0.0
ub_a12 = 10.0
lb_a21 = -10.0
ub_a21 = 0.0
a11 = Parameter("a11", 1, 0.0)
a12 = Parameter("a12", 1, lb_a12, ub_a12)
a21 = Parameter("a21", 1, lb_a21, ub_a21)
a22 = Parameter("a22", 1, ub=0.0)
params = [a11, a12, a21, a22]
M = Model("lds_2D", params)
M.set_eps(linear2D_freq)
q_theta, opt_data, epi_path, failed = M.epi(
mu, num_iters=100, K=1, save_movie_data=True, log_rate=10,
)
z = q_theta(50)
g = q_theta.plot_dist(z)
M.epi_opt_movie(epi_path)
params = [a11, a12, a21, a22]
# should load from prev epi
M = Model("lds_2D", params)
M.set_eps(linear2D_freq)
q_theta, opt_data, epi_path, failed = M.epi(
mu, num_iters=100, K=1, save_movie_data=True
)
print("epi_path", epi_path)
epi_df = M.get_epi_df()
epi_df_row = epi_df[epi_df["iteration"] == 100].iloc[0]
q_theta = M.get_epi_dist(epi_df_row)
opt_data_filename = os.path.join(epi_path, "opt_data.csv")
M.set_eps(linear2D_freq)
q_theta, opt_data, epi_path, failed = M.epi(
mu, num_iters=100, K=1, save_movie_data=True, log_rate=10,
)
opt_data_cols = ["k", "iteration", "H", "cost", "converged"] + [
"R%d" % i for i in range(1, M.m + 1)
]
for x, y in zip(opt_data.columns, opt_data_cols):
assert x == y
assert q_theta is not None
z = q_theta(1000)
log_q_z = q_theta.log_prob(z)
assert np.sum(z[:, 0] < 0.0) == 0
assert np.sum(z[:, 1] < lb_a12) == 0
assert np.sum(z[:, 1] > ub_a12) == 0
assert np.sum(z[:, 2] < lb_a21) == 0
assert np.sum(z[:, 2] > ub_a21) == 0
assert np.sum(z[:, 3] > 0.0) == 0
assert np.sum(1 - np.isfinite(z)) == 0
# Intentionally swap order in list to insure proper handling.
params = [a22, a21, a12, a11]
M = Model("lds", params)
M.set_eps(linear2D_freq)
q_theta, opt_data, epi_path, _ = M.epi(
mu,
K=2,
num_iters=100,
stop_early=True,
verbose=True,
save_movie_data=True,
log_rate=10,
)
M.epi_opt_movie(epi_path)
z = q_theta(1000)
log_q_z = q_theta.log_prob(z)
assert np.sum(z[:, 0] < 0.0) == 0
assert np.sum(z[:, 1] < lb_a12) == 0
assert np.sum(z[:, 1] > ub_a12) == 0
assert np.sum(z[:, 2] < lb_a21) == 0
assert np.sum(z[:, 2] > ub_a21) == 0
assert np.sum(z[:, 3] > 0.0) == 0
assert np.sum(1 - np.isfinite(z)) == 0
print("DOING ABC NOW")
# Need finite support for ABC
a11 = Parameter("a11", 1, -10.0, 10.0)
a12 = Parameter("a12", 1, -10.0, 10.0)
a21 = Parameter("a21", 1, -10.0, 10.0)
a22 = Parameter("a22", 1, -10.0, 10.0)
params = [a11, a12, a21, a22]
M = Model("lds_2D", params)
M.set_eps(linear2D_freq)
init_type = "abc"
init_params = {"num_keep": 50, "mean": mu[:2], "std": np.sqrt(mu[2:])}
q_theta, opt_data, epi_path, failed = M.epi(
mu,
num_iters=100,
K=1,
init_type=init_type,
init_params=init_params,
save_movie_data=True,
log_rate=10,
)
params = [a11, a12, a21, a22]
M = Model("lds2", params)
M.set_eps(linear2D_freq)
# This should cause opt to fail with nan since c0=1e20 is too high.
q_theta, opt_data, epi_path, _ = M.epi(
mu,
K=3,
num_iters=1000,
c0=1e20,
stop_early=True,
verbose=True,
save_movie_data=False,
log_rate=10,
)
with raises(IOError):
M.epi_opt_movie(epi_path)
for x, y in zip(opt_data.columns, opt_data_cols):
assert x == y
with raises(ValueError):
def bad_f(a11, a12, a21, a22):
return tf.expand_dims(a11 + a12 + a21 + a22, 0)
M.set_eps(bad_f)
params = [a11, a12, a21, a22]
M = Model("lds2", params)
init_params = {"mu": 2 * np.zeros((4,)), "Sigma": np.eye(4)}
nf = NormalizingFlow("autoregressive", 4, 1, 2, 10)
al_hps = AugLagHPs()
epi_path, exists = M.get_epi_path(init_params, nf, mu, al_hps, eps_name="foo")
assert not exists
return None