def make_nmf(fix_atoms=False):
dim_k, dim_f, dim_t = 5, 20, 100
atoms_kf = npr.dirichlet(np.ones(dim_f) * 0.9, size=dim_k)
activations_tk = npr.gamma(shape=0.6, scale=200, size=(dim_t, dim_k))
observations_tf = np.einsum("tk,kf->tf", activations_tk, atoms_kf)
observations_tf += npr.randn(dim_t, dim_f) * 1e-4
leaf_kf = wtf.LeafDirichlet(
name="atoms",
index_id="kf",
norm_axis=(1,),
tensor=atoms_kf if fix_atoms else wtfu.normalize(npr.rand(dim_k, dim_f), (1,)),
update_period=0 if fix_atoms else 1,
prior_shape=1,
)
leaf_tk = wtf.LeafGamma(
name="activations",
index_id="tk",
tensor=np.ones_like(activations_tk),
prior_rate=1e-5,
prior_shape=1,
)
mul_tf = wtf.Multiplier(name="multiplier", index_id="tf")
mul_tf.new_parents(leaf_kf, leaf_tk)
obs_tf = wtf.PosObserver(name="observer", index_id="tf", tensor=observations_tf)
mul_tf.new_child(obs_tf)
root = wtf.Root(name="nmf")
obs_tf.new_child(root)
return root
def make_smooth_activation_nmf():
"""
NMF with overlapping activation (each time coefficient spreads out on
both sides). We uses this model in order to test strides_for_child
"""
_, leaf_ktw, mul_kwt, mul_tf, obs_tf = _aux_smooth_activation_nmf()
dim_k, _, dim_w = leaf_ktw.tensor.shape
dim_t = obs_tf.tensor.shape[0]
leaf_w = wtf.LeafDirichlet(
name="adder", index_id="w", norm_axis=(), tensor=np.ones(dim_w), update_period=0
)
mul_tk = wtf.Multiplier(name="activations", index_id="tk")
shape = (dim_k, dim_w, dim_t + dim_w - 1)
strides_for_child = (
[
8,
]
+ list(np.cumprod(shape[:0:-1]) * 8)
)[::-1]
strides_for_child[-2] += 8
mul_kwt.new_child(
mul_tk,
shape_for_child=(dim_k, dim_w, dim_t),
strides_for_child=strides_for_child,
)
leaf_w.new_child(mul_tk)
mul_tf.new_parent(mul_tk)
root = wtf.Root(name="smooth_activation_nmf")
obs_tf.new_child(root)
return root
def _aux_smooth_activation_nmf():
dim_k, dim_f, dim_t = 5, 20, 100
dim_w = 9
atoms_kf = npr.dirichlet(np.ones(dim_f) * 0.9, size=dim_k)
activations_tk = npr.gamma(shape=0.6, scale=200, size=(dim_t + dim_w - 1, dim_k))
spread_tkw = npr.dirichlet(3 * np.ones(dim_w), size=(dim_t + dim_w - 1, dim_k))
act_tkw = np.einsum("tk,tkw->tkw", activations_tk, spread_tkw)
activations_tk = np.zeros((dim_t, dim_k))
for ww in range(dim_w):
activations_tk += act_tkw[ww : dim_t + ww, :, ww]
observations_tf = np.einsum("tk,kf->tf", activations_tk, atoms_kf)
observations_tf += npr.randn(dim_t, dim_f) * 1e-4
leaf_kf = wtf.LeafDirichlet(
name="atoms",
index_id="kf",
norm_axis=(1,),
tensor=wtfu.normalize(npr.rand(dim_k, dim_f), (1,)),
prior_shape=1,
)
leaf_kt = wtf.LeafGamma(
name="impulse",
index_id="kt",
tensor=np.ones((dim_k, dim_t + dim_w - 1)),
prior_shape=1,
prior_rate=1e-5,
)
leaf_ktw = wtf.LeafDirichlet(
name="spreading",
index_id="ktw",
norm_axis=(2,),
tensor=np.ones((dim_k, dim_t + dim_w - 1, dim_w)) / dim_w,
prior_shape=10,
)
mul_kwt = wtf.Multiplier(name="mul_kwt", index_id="kwt")
mul_kwt.new_parents(leaf_kt, leaf_ktw)
mul_tf = wtf.Multiplier(name="reconstruction", index_id="tf")
mul_tf.new_parent(leaf_kf)
obs_tf = wtf.PosObserver(name="observer", index_id="tf", tensor=observations_tf)
mul_tf.new_child(obs_tf)
return leaf_kt, leaf_ktw, mul_kwt, mul_tf, obs_tf
def make_deconv_tree():
dim_x0, dim_y0, dim_x1, dim_y1 = 20, 20, 3, 3
dim_x = dim_x0 - dim_x1 + 1
dim_y = dim_y0 - dim_y1 + 1
dim_k = 2
kernel_kyx = np.zeros((dim_k, dim_y1, dim_x1))
kernel_kyx[0, [1, 0, 1, 2, 1], [0, 1, 1, 1, 2]] = 1
kernel_kyx[1, [0, 0, 1, 1, 2, 2, 2], [0, 2, 0, 2, 0, 1, 2]] = 1
kernel_kyx /= kernel_kyx.sum((1, 2), keepdims=True)
impulse_kyx = npr.gamma(shape=0.08, scale=200, size=(dim_k, dim_y0, dim_x0))
impulse_kyx[impulse_kyx < 200] = 0
impulse_kyx[impulse_kyx >= 200] = 200
image_yx = np.zeros((dim_y, dim_x))
for kk in range(dim_k):
image_yx += convolve2d(impulse_kyx[kk], kernel_kyx[kk], mode="valid")
leaf_kernel = wtf.LeafDirichlet(
name="kernel",
index_id=("k", "j", "i"),
norm_axis=(1, 2),
tensor=np.ones((dim_k, dim_y1, dim_x1)),
prior_shape=100 + 1e-4 * npr.rand(dim_k, dim_y1, dim_x1),
init_type="prior",
)
leaf_impulse = wtf.LeafGamma(
name="impulse",
index_id=("k", "y", "x"),
tensor=np.ones((dim_k, dim_y0, dim_x0)),
prior_shape=1,
prior_rate=0.0001,
)
mul_image = wtf.Multiplier(
name="reconstruction", index_id="yx", conv_idx_ids=("y", "x")
)
leaf_kernel.new_child(mul_image, index_id_for_child=("k", "y", "x"))
leaf_impulse.new_child(mul_image, index_id_for_child=("k", "y", "x"))
obs = wtf.PosObserver(
name="image",
index_id="yx",
tensor=image_yx,
drawings_max=200 * dim_x * dim_y,
drawings_step=10,
)
mul_image.new_child(obs)
root = wtf.Root(name="root")
obs.new_child(root)
return root
def make_sparse_nmf2(prior_rate=0.001, obs=None):
"""
NMF with l1/l2 sparse constraint on atoms
"""
dim_f, dim_t, dim_k = 2, 100, 2
# gt_kf = npr.dirichlet(np.ones(dim_f), size=dim_k)
gt_kf = np.array([[4.0, 1.0], [4.0, 3.0]])
gt_kf /= gt_kf.sum(1, keepdims=True)
gt_tk = npr.gamma(shape=0.3, scale=100, size=(dim_t, dim_k))
gt_tf = np.dot(gt_tk, gt_kf)
gt_tf += npr.rand(dim_t, dim_f)
if obs is not None:
gt_tf = obs
leaf_kf = wtf.LeafGammaNorm(
name="atoms",
index_id="kf",
tensor=np.ones((dim_k, dim_f)),
l2_norm_axis=(1,),
prior_rate=prior_rate,
prior_shape=1 + 1e-4 * npr.rand(dim_k, dim_k),
)
leaf_kt = wtf.LeafDirichlet(
name="activations",
index_id="kt",
norm_axis=(1,),
tensor=np.ones((dim_k, dim_t)) / dim_t,
prior_shape=1,
)
mul_tf = wtf.Multiplier(name="reconstruction", index_id="tf")
mul_tf.new_parents(leaf_kt, leaf_kf)
obs_tf = wtf.PosObserver(name="observations", index_id="tf", tensor=gt_tf)
mul_tf.new_child(obs_tf)
root = wtf.Root(name="root", verbose_iter=50, cost_computation_iter=10)
obs_tf.new_child(root)
return root
def make_test_tree():
"""
A toy wonterfact tree is designed in order to test the maximum
possibilities.
"""
dim_k, dim_s, dim_t, dim_f, dim_c = 3, 2, 10, 4, 2
leaf_energy = wtf.LeafGamma(
name="leaf_energy",
index_id="",
tensor=np.array(1),
prior_shape=1.1,
prior_rate=0.01,
)
leaf_k = wtf.LeafDirichlet(
name="leaf_k",
index_id="k",
norm_axis=(0, ),
tensor=normalize(npr.rand(dim_k), 0),
prior_shape=1,
)
mul_k = wtf.Multiplier(name="mul_k", index_id=("k", ))
mul_k.new_parents(leaf_energy, leaf_k)
leaf_kts_0 = wtf.LeafDirichlet(
name="leaf_kts_0",
index_id="kts",
norm_axis=(1, 2),
tensor=normalize(npr.rand(dim_k, dim_t, dim_s), (1, 2)),
prior_shape=1,
)
mul_kts = wtf.Multiplier(name="mul_kts", index_id="kts")
mul_kts.new_parents(leaf_kts_0, mul_k)
leaf_kts_1 = wtf.LeafGamma(
name="leaf_kts_1",
index_id="kts",
tensor=npr.rand(dim_k, dim_t, dim_s),
prior_shape=1,
prior_rate=0.001,
)
mult_ktsc = wtf.Multiplexer(name="mult_ktsc", index_id="ktsc")
mult_ktsc.new_parents(leaf_kts_1, mul_kts)
# no update for this leaf
# it acts like an Adder (no normalization)
leaf_c = wtf.LeafDirichlet(
name="leaf_c",
norm_axis=(),
index_id="c",
tensor=np.ones(dim_c),
update_period=0,
)
mul_kts_1 = wtf.Multiplier(name="mul_kts_1", index_id="kts")
mul_kts_1.new_parents(mult_ktsc, leaf_c)
# two updates every 3 iterations
leaf_kf = wtf.LeafDirichlet(
name="leaf_kf",
index_id="kf",
norm_axis=(1, ),
tensor=normalize(npr.rand(dim_k, dim_f), 1),
prior_shape=1,
update_period=3,
update_succ=2,
update_offset=0,
)
mul_tfs = wtf.Multiplier(name="mul_tfs", index_id="tfs")
mul_tfs.new_parents(mul_kts_1, leaf_kf)
obs_tf = wtf.RealObserver(name="obs_tf",
index_id="tf",
tensor=100 * npr.randn(dim_t, dim_f))
mul_tfs.new_child(obs_tf)
root = wtf.Root(
name="root",
cost_computation_iter=1,
stop_estim_threshold=0,
update_type="regular",
inference_mode="EM",
verbose_iter=10,
)
obs_tf.new_child(root)
leaf_k_1 = wtf.LeafGamma(
name="leaf_k_1",
index_id="k",
tensor=npr.rand(dim_k),
prior_shape=1,
prior_rate=0.001,
)
# one update every 3 iterations
leaf_kt = wtf.LeafDirichlet(
name="leaf_kt",
index_id="kt",
norm_axis=(1, ),
tensor=normalize(np.ones((dim_k, dim_t)), 1),
prior_shape=1,
update_period=3,
update_succ=1,
update_offset=2,
)
mul_kt = wtf.Multiplier(name="mul_kt", index_id="kt")
mul_kt.new_parents(leaf_k_1, leaf_kt)
mul_tf = wtf.Multiplier(name="mul_tf", index_id="tf")
mul_tf.new_parents(leaf_kf, mul_kt)
obs_tf_2 = wtf.PosObserver(name="obs_tf_2",
index_id="tf",
tensor=100 * npr.rand(dim_t, dim_f))
mul_tf.new_child(obs_tf_2)
obs_tf_2.new_child(root)
root.dim_k, root.dim_s, root.dim_t, root.dim_f, root.dim_c = (
dim_k,
dim_s,
dim_t,
dim_f,
dim_c,
)
return root
def make_sparse_nmf3(prior_rate=0.001, obs=None):
"""
NMF with approximation of l2 norm for atoms
"""
dim_f, dim_t, dim_k, dim_a = 2, 100, 2, 2
gt_kf = npr.dirichlet(np.ones(dim_f), size=dim_k)
# gt_kf = np.array([[4.0, 1.0], [4.0, 3.0]])
gt_kf /= gt_kf.sum(1, keepdims=True)
gt_tk = npr.gamma(shape=0.3, scale=100, size=(dim_t, dim_k))
gt_tf = np.dot(gt_tk, gt_kf)
gt_tf += npr.rand(dim_t, dim_f)
if obs is not None:
gt_tf = obs
leaf_k = wtf.LeafGamma(
name="atoms_energy",
index_id="k",
tensor=np.ones((dim_k)),
prior_shape=1,
prior_rate=prior_rate,
)
leaf_kf = wtf.LeafDirichlet(
name="atoms_init",
index_id="kf",
norm_axis=(1,),
tensor=wtfu.normalize(npr.rand(dim_k, dim_f), (1,)),
prior_shape=1,
)
mul_kf = wtf.Multiplier(index_id="kf")
mul_kf.new_parents(leaf_kf, leaf_k)
mul_kfg = wtf.Multiplier(index_id="kfg")
mul_kf.new_child(mul_kfg)
leaf_kf.new_child(mul_kfg, index_id_for_child="kg")
leaf_c = wtf.LeafDirichlet(
index_id="c", norm_axis=(0,), tensor=np.array([0.5, 0.5]), update_period=0
)
mul_ckf = wtf.Multiplier(index_id="ckf")
leaf_c.new_child(mul_ckf)
test_arr = npr.rand(2, dim_f, dim_f)
strides = (test_arr.strides[0],) + np.diag(test_arr[0]).strides
mul_kfg.new_child(
mul_ckf,
index_id_for_child="kf",
shape_for_child=(dim_k, dim_f),
strides_for_child=strides,
)
leaf_g = wtf.LeafDirichlet(
index_id="g", norm_axis=(), tensor=np.ones(dim_f - 1), update_period=0
)
mul_kf2 = wtf.Multiplier(index_id="kf")
leaf_g.new_child(mul_kf2)
mask = np.logical_not(np.eye(dim_f, dtype=bool))
mul_kfg.new_child(
mul_kf2,
slice_for_child=[slice(None), mask],
shape_for_child=(dim_k, dim_f, dim_f - 1),
)
add_kf = wtf.Adder(name="atoms", index_id="kf")
mul_kf2.new_child(add_kf)
mul_ckf.new_child(add_kf, index_id_for_child="kf", slice_for_child=(0, Ellipsis))
# leaf_ka = wtf.LeafDirichlet(
# name="angle",
# index_id="ka",
# norm_axis=(1,),
# tensor=np.ones((dim_k, dim_a)),
# l2_norm_axis=(1,),
# prior_shape=1 + 1e-4 * npr.rand(dim_k, dim_a),
# )
# mul_ka = wtf.Multiplier(index_id="ka")
# mul_ka.new_parents(leaf_k, leaf_ka)
# mul_kab = wtf.Multiplier(index_id="kab")
# mul_ka.new_child(mul_kab)
# leaf_ka.new_child(mul_kab, index_id_for_child="kb")
# leaf_abm = wtf.LeafDirichlet(
# index_id="abm",
# norm_axis=(2,),
# tensor=np.array([[[1, 0, 0], [0, 1, 0]], [[0, 1, 0], [0, 0, 1]]]),
# update_period=0,
# )
# mul_km = wtf.Multiplier(index_id="km")
# mul_km.new_parents(leaf_abm, mul_kab)
# leaf_mnf = wtf.LeafDirichlet(
# name="basis",
# index_id="mnf",
# norm_axis=(1, 2),
# tensor=np.array(
# [[[0, 0.5], [0.5, 0]], [[0.5, 0.5], [0, 0]], [[0.5, 0], [0, 0.5]]]
# ),
# update_period=0,
# )
# mul_nkf = wtf.Multiplier(index_id="nkf", name="atoms")
# mul_nkf.new_parents(leaf_mnf, mul_km)
leaf_kt = wtf.LeafDirichlet(
name="activations",
index_id="kt",
norm_axis=(1,),
tensor=np.ones((dim_k, dim_t)) / dim_t,
prior_shape=1,
)
mul_tf = wtf.Multiplier(name="reconstruction", index_id="tf")
leaf_kt.new_child(mul_tf)
# mul_nkf.new_child(mul_tf, index_id_for_child="kf", slice_for_child=[0, Ellipsis])
add_kf.new_child(mul_tf)
obs_tf = wtf.PosObserver(name="observations", index_id="tf", tensor=gt_tf)
mul_tf.new_child(obs_tf)
root = wtf.Root(name="root", verbose_iter=50, cost_computation_iter=10)
obs_tf.new_child(root)
return root
def make_sparse_nmf(prior_rate=0.001, obs=None, atoms=None):
"""
NMF with minimization of \\sum_{k != k'} P(k|t)P(k'|t)E(t) where P(k|t) are
the activations and E(t) total energy at time t
"""
dim_f, dim_t, dim_k = 2, 100, 2
# gt_kf = npr.dirichlet(np.ones(dim_f), size=dim_k)
gt_kf = np.array([[4.0, 1.0], [4.0, 3.0]])
gt_kf /= gt_kf.sum(1, keepdims=True)
gt_tk = npr.gamma(shape=0.3, scale=100, size=(dim_t, dim_k))
gt_tf = np.dot(gt_tk, gt_kf)
gt_tf += npr.rand(dim_t, dim_f)
if obs is not None:
gt_tf = obs
leaf_t = wtf.LeafGamma(
name="time_energy",
index_id="t",
tensor=np.ones(dim_t),
prior_rate=prior_rate,
prior_shape=1,
)
leaf_tk = wtf.LeafDirichlet(
name="activations",
index_id="tk",
norm_axis=(1,),
tensor=np.ones((dim_t, dim_k)) / dim_k,
prior_shape=1,
)
mul_tk = wtf.Multiplier(index_id="tk")
mul_tk.new_parents(leaf_t, leaf_tk)
mul_tkl = wtf.Multiplier(name="activations_square", index_id="tkl")
leaf_tk.new_child(mul_tkl, index_id_for_child="tl")
mul_tk.new_child(mul_tkl)
if atoms is None:
atoms = np.ones((dim_k, dim_f)) / dim_f
update_period = 1
else:
update_period = 0
leaf_kf = wtf.LeafDirichlet(
name="atoms",
index_id="kf",
norm_axis=(1,),
tensor=atoms,
prior_shape=1,
update_period=update_period,
)
mul_tf = wtf.Multiplier(name="reconstruction", index_id="tf")
leaf_kf.new_child(mul_tf)
test_arr = npr.rand(2, dim_k, dim_k)
strides = (test_arr.strides[0],) + np.diag(test_arr[0]).strides
mul_tkl.new_child(
mul_tf,
shape_for_child=(dim_t, dim_k),
strides_for_child=strides,
index_id_for_child="tk",
)
obs_tf = wtf.PosObserver(name="observations", index_id="tf", tensor=gt_tf)
mul_tf.new_child(obs_tf)
root = wtf.Root(name="root", verbose_iter=50, cost_computation_iter=10)
obs_tf.new_child(root)
return root