from skued import diffread, bragg_peaks_persistence, bragg_peaks
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Circle

im = diffread("data/ab_initio_monolayer_mos2.tif")

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6, 3))

peaks = np.array(bragg_peaks(im, mask=np.ones_like(im, dtype=bool))).reshape(-1,2)
peaks_persistence, _, _, _ = bragg_peaks_persistence(im, prominence=0.1)
for pp, image, ax in zip([peaks, peaks_persistence], [im, im], [ax1, ax2]):
        ax.imshow(im, origin='lower', interpolation='bicubic', vmin=-1, vmax=1)
        for r, c in pp:
                ax.add_patch(Circle(xy=(r, c), radius=5, ec="r", fc="none"))
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
plt.tight_layout()
plt.show()