import matplotlib.pyplot as plt
import h5py
import sys
import numpy as np

dataset = sys.argv[1]
data = h5py.File(dataset, "r")["data"]
idxs = np.random.randint(data.shape[0], size=100)
vmax = np.abs(data).max()
axes = []
D = data.shape[1]
R = int(np.sqrt(D))
for i in range(100):
    ax = plt.subplot(10, 10, i + 1)
    axes.append(ax)
    im = ax.imshow(data[i].reshape(R,R), vmin=-vmax, vmax=vmax, cmap="coolwarm")
    plt.axis("off")
plt.colorbar(im, ax=axes)
plt.show()
