<jupyter_start><jupyter_text>Generate a count matrix of random noiseWe want our factor model to tell us that this is just noise<jupyter_code>N = 50000
D = 30
P = 4

# Test taking in from tf.dataset, don't pre-batch
data = tf.data.Dataset.from_tensor_slices(
    {
        'counts': np.random.poisson(1.0, size=(N, D)),
        'indices': np.arange(N),
    })

data = data.batch(1000)<jupyter_output><empty_output><jupyter_text>Currently, the input data has to be in this format, as a tensorflow BatchDataset object with the given key/values.  Set up the inference object<jupyter_code># strategy = tf.distribute.MirroredStrategy()
strategy = None
factor = PoissonMatrixFactorization(
    data, latent_dim=P, strategy=strategy,
    u_tau_scale=1.0/np.sqrt(D*N),
    dtype=tf.float64)
# Test to make sure sampling works<jupyter_output>Looping through the entire dataset once to get some stats
Feature dim: 30 -> Latent dim 4<jupyter_text>Inference<jupyter_code>losses = factor.calibrate_advi(
    num_epochs=200, rel_tol=1e-4, learning_rate=.05)<jupyter_output>Initial loss: 44.13152280855831
Epoch 1: average-batch loss: 43.353231421619064 last batch loss: 43.000628571823775
Saved a checkpoint: ./.tf_ckpts/3c264df9-4e73-403f-adf3-1a8fee143964/3c264df9-4e73-403f-adf3-1a8fee143964-1
Epoch 2: average-batch loss: 42.56571656362709 last batch loss: 42.329743119167155
Saved a checkpoint: ./.tf_ckpts/3c264df9-4e73-403f-adf3-1a8fee143964/3c264df9-4e73-403f-adf3-1a8fee143964-2
Epoch 3: average-batch loss: 41.956856755999226 last batch loss: 41.71911894726435
Saved a checkpoint: ./.tf_ckpts/3c264df9-4e73-403f-adf3-1a8fee143964/3c264df9-4e73-403f-adf3-1a8fee143964-3
Epoch 4: average-batch loss: 41.385839702186715 last batch loss: 41.23389101231501
Saved a checkpoint: ./.tf_ckpts/3c264df9-4e73-403f-adf3-1a8fee143964/3c264df9-4e73-403f-adf3-1a8fee143964-4
Epoch 5: average-batch loss: 41.02195755552638 last batch loss: 40.979339612180816
Saved a checkpoint: ./.tf_ckpts/3c264df9-4e73-403f-adf3-1a8fee143964/3c264df9-4e73-403f-adf3-1a8fee143964-5
Epoch 6: ave[...]<jupyter_text>Compute the WAIC<jupyter_code>waic = factor.waic()
print(waic)<jupyter_output>{'waic': 345148.4950137764, 'se': 1761.722935996936, 'lppd': -37090.95152008469, 'pwaic': 135483.2959868035}<jupyter_text>Plot the factor encoding matrix<jupyter_code>surrogate_samples = factor.surrogate_distribution.sample(1000)
if 's' in surrogate_samples.keys():
    weights = surrogate_samples['s']/tf.reduce_sum(surrogate_samples['s'],-2,keepdims=True)
    intercept_data = az.convert_to_inference_data(
        {
            r"$\varphi_i$": 
                (tf.squeeze(surrogate_samples['w'])*weights[:,-1,:]*factor.eta_i).numpy().T})
else:
    intercept_data = az.convert_to_inference_data(
    {
        r"$\varphi_i$": 
            (tf.squeeze(surrogate_samples['w'])*factor.eta_i).numpy().T})
fig, ax = plt.subplots(figsize=(10,4))

plt.imshow(factor.encoding_matrix().numpy().T, vmin=0, cmap="Blues")
plt.xticks(np.arange(D))
ax.set_xlabel("item")
ax.set_ylabel("dimension")
plt.tick_params(axis='both', which='major', labelsize=10, labelbottom = False, bottom=False, top = True, labeltop=True)
plt.yticks(np.arange(P))
plt.colorbar(orientation = "horizontal")
plt.show()

fig, ax = plt.subplots(1,2, figsize=(14,8))

pcm = ax[0].imshow(factor.encoding_matrix().numpy()[::-1,:], vmin=0, cmap="Blues")
ax[0].set_yticks(np.arange(D))
ax[0].set_yticklabels(np.arange(D))
ax[0].set_ylabel("item")
ax[0].set_xlabel("factor dimension")
ax[0].set_xticks(np.arange(P))
ax[0].set_xticklabels(np.arange(P))

fig.colorbar(pcm, ax=ax[0], orientation = "vertical")
az.plot_forest(intercept_data, ax=ax[1])
ax[1].set_xlabel("background rate")
ax[1].set_ylim((-0.014,.466))
ax[1].set_title("65% and 95% CI")
ax[1].axvline(1.0, linestyle='dashed', color="black")
#plt.savefig('noise_factorization_sepmf.pdf', bbox_inches='tight')
plt.show()<jupyter_output><empty_output><jupyter_text>The values in the factor encoding matrix as near zero. There is a gradient because the model induces a dimensional decay to remove symmetries - this encourages the leading dimensions to hold more of the signal.<jupyter_code>fig, ax = plt.subplots(1,2, figsize=(14,8))

pcm = ax[0].imshow(factor.decoding_matrix().numpy().T[::-1,:], vmin=0, cmap="Blues")
ax[0].set_yticks(np.arange(D))
ax[0].set_yticklabels(np.arange(D))
ax[0].set_ylabel("item")
ax[0].set_xlabel("factor dimension")
ax[0].set_xticks(np.arange(P))
ax[0].set_xticklabels(np.arange(P))

fig.colorbar(pcm, ax=ax[0], orientation = "vertical")
az.plot_forest(intercept_data, ax=ax[1])
ax[1].set_xlabel("background rate")
ax[1].set_ylim((-0.014,.466))
ax[1].set_title("65% and 95% CI")
ax[1].axvline(1.0, linestyle='dashed', color="black")
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1549738219895288, 'fraction_numerical': 0.11790575916230367, 'mean_word_length': 5.972262773722628, 'pattern_counts': {'":': 2, '<': 22, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '603603', 'license': 'MIT', 'n_tokens_mistral': 2271, 'n_tokens_neox': 1949, 'n_words': 322, 'path': 'notebooks/factorizing_random_noise.ipynb', 'repo_name': 'mederrata/spmf', 'size': '4775'}