[mngseg] add gaussian / coherent PS segementation tool

mngseg.py

@@ -0,0 +1,275 @@
+from pywavan import powspec, nb_scale
+from astropy.io import fits
+import numpy as np
+from matplotlib import pyplot as plt
+import aplpy
+
+
+def make_images(im, wt, M, meanim, label):
+    """ show the Gaussian and coherent part of the image """
+
+    total = np.sum(wt[:M, :, :], axis=0).real + meanim
+    coherent = np.sum(wt[M : 2 * M, :, :], axis=0).real + meanim
+    Gaussian = np.sum(wt[2 * M : 3 * M, :, :], axis=0).real + meanim
+
+    fig_all = plt.figure(1, figsize=(16, 4))
+
+    # original image
+    fig = aplpy.FITSFigure(fits.PrimaryHDU(im), figure=fig_all, subplot=(1, 4, 1))
+    fig.show_colorscale(cmap="cubehelix")
+    fig.add_colorbar()
+    fig.axis_labels.hide()
+    fig.tick_labels.hide()
+    fig.set_title("Original")
+
+    # wavelet image total (should be same as original image)
+    fig = aplpy.FITSFigure(fits.PrimaryHDU(total), figure=fig_all, subplot=(1, 4, 2))
+    fig.show_colorscale(cmap="cubehelix")
+    fig.add_colorbar()
+    fig.axis_labels.hide()
+    fig.tick_labels.hide()
+    fig.set_title("wavelet")
+
+    # gaussian component
+    fig = aplpy.FITSFigure(fits.PrimaryHDU(Gaussian), figure=fig_all, subplot=(1, 4, 3))
+    fig.show_colorscale(cmap="cubehelix")
+    fig.add_colorbar()
+    fig.axis_labels.hide()
+    fig.tick_labels.hide()
+    fig.set_title("Gaussian")
+
+    # coherent component
+    fig = aplpy.FITSFigure(fits.PrimaryHDU(coherent), figure=fig_all, subplot=(1, 4, 4))
+    fig.show_colorscale(cmap="cubehelix")
+    fig.add_colorbar()
+    fig.axis_labels.hide()
+    fig.tick_labels.hide()
+    fig.set_title("Coherent")
+
+    plt.tight_layout()
+    plt.savefig("reconstructed_image_{}.png".format(label))
+    plt.close()
+
+
+def scale_images(thingy, M, label, scale=14, mode="wt"):
+    """ visualize wt or S11a for a specific scale. Remark S11a = wt^2"""
+    total = thingy[scale, :, :].real
+    coherent = thingy[M + scale, :, :].real
+    Gaussian = thingy[2 * M + scale, :, :].real
+
+    # make images
+    fig_all = plt.figure(1, figsize=(12, 4))
+
+    # wavelet image on scale
+    fig = aplpy.FITSFigure(fits.PrimaryHDU(total), figure=fig_all, subplot=(1, 3, 1))
+    limit = max(np.max(total), abs(np.min(total)))
+    fig.show_colorscale(cmap="PiYG", vmin=-limit, vmax=limit)
+    fig.add_colorbar()
+    fig.axis_labels.hide()
+    fig.tick_labels.hide()
+    fig.set_title("wavelet")
+
+    # gaussian component
+    fig = aplpy.FITSFigure(fits.PrimaryHDU(Gaussian), figure=fig_all, subplot=(1, 3, 2))
+    limit = max(np.max(Gaussian), abs(np.min(Gaussian)))
+    fig.show_colorscale(cmap="PiYG", vmin=-limit, vmax=limit)
+    fig.add_colorbar()
+    fig.axis_labels.hide()
+    fig.tick_labels.hide()
+    fig.set_title("Gaussian")
+
+    # coherent component
+    fig = aplpy.FITSFigure(fits.PrimaryHDU(coherent), figure=fig_all, subplot=(1, 3, 3))
+    limit = max(np.max(coherent), abs(np.min(coherent)))
+    fig.show_colorscale(cmap="PiYG", vmin=-limit, vmax=limit)
+    fig.add_colorbar()
+    fig.axis_labels.hide()
+    fig.tick_labels.hide()
+    fig.set_title("Coherent")
+
+    plt.tight_layout()
+    plt.savefig("imag_{}_scale{}_{}.png".format(mode, scale, label))
+    plt.close()
+
+
+def plot_each_scale(S11a, wav_k, q, label, coherent=False, reso=1):
+    """ plot histogram at a certain scale """
+    nsize = len(S11a[0, 0, :])
+    M = len(q)
+    for scl in range(0, M):
+        plt.figure(figsize=(6, 6))
+        # determine bins (large scales should have less bins)
+        nbins = np.int(nsize ** 2.0 * (wav_k[scl] * reso) ** 2.0)
+        nbins = max(9, nbins)
+        nbins = min(500, nbins)
+        # calc histogram gaussian component w.r.t. its mean value (easier to compare)
+        intermit = (S11a[2 * M + scl, :, :]) / np.mean(S11a[2 * M + scl, :, :])
+        histo, edges = np.histogram(intermit, density=True, bins=nbins)
+        plt.hist(
+            (edges[:-1] + edges[1:]) / 2.0,
+            bins=edges,
+            weights=histo,
+            alpha=0.5,
+            label="Gaussian",
+            color="red",
+        )
+        # calc histogram coherent component
+        if coherent and (np.mean(S11a[M + scl, :, :]) > 0):
+            intermit_C = (S11a[M + scl, :, :]) / np.mean(S11a[M + scl, :, :])
+            histo_C, edges_C = np.histogram(intermit_C, density=True, bins=nbins)
+            plt.hist(
+                (edges_C[:-1] + edges_C[1:]) / 2.0,
+                bins=edges_C,
+                weights=histo_C,
+                alpha=0.5,
+                label="coherent",
+                color="blue",
+            )
+        else:
+            histo_C = []
+        # plot mean
+        plt.plot([1, 1], [0, 5])
+        plt.xlabel(r"$I/\langle I \rangle$")
+        plt.ylabel("PDF")
+        plt.title(
+            "scale {}, k= {}, l={} pc".format(
+                scl, np.str(wav_k[scl])[:7], 1000.0 / nsize / wav_k[scl]
+            )
+        )
+        plt.legend()
+        # avoid first bin dominating the range
+        gauss_max = 0
+        if len(histo) > 1:
+            gauss_max = max(histo[1:])
+        coh_max = 0
+        if coherent & len(histo_C) > 1:
+            coh_max = max(histo_C[1:])
+        if max(gauss_max, coh_max) > 0:
+            plt.ylim(0, max(gauss_max, coh_max))
+        plt.xlim(0, 4)
+        plt.ylim(0, 1.0)
+        plt.tight_layout()
+        plt.savefig("S11a_scale{}_q_{}_{}.png".format(scl, q[scl], label))
+        plt.close()
+
+
+def plot_components_power_spectrum(
+    tab_k, wav_k, spec_k, S1a, label="", fit_min=7, fit_max=15, nsize=1024, lvlmin=8
+):
+
+    # plot power spectra
+    plt.figure(figsize=(5, 5))
+    plt.plot(tab_k, spec_k, color="black", label="Fourier PS", linewidth=1.5)
+    plt.plot(wav_k, S1a[0, :], "s", color="black", label="Wavelet PS", markersize=6)
+    plt.plot(wav_k, S1a[1, :], "D", color="blue", label="Coherent PS", markersize=4)
+    plt.plot(wav_k, S1a[2, :], "^", color="red", label="Gaussian PS", markersize=5)
+
+    scl_min = fit_min
+    scl_max = fit_max + 1
+
+    # power law fit to gaussian component
+    coefw, cov = np.polyfit(
+        np.log(wav_k[scl_min:scl_max]), np.log(S1a[2, scl_min:scl_max]), deg=1, cov=True
+    )
+    yfitw = np.exp(coefw[1]) * wav_k[scl_min:scl_max] ** coefw[0]
+    plt.plot(wav_k[scl_min:scl_max], yfitw, "-", color=("red"), linewidth=4, alpha=0.5)
+    print("Gaussian Power law", coefw[0])
+
+    # Coherent power law fit
+    coefcw, cov = np.polyfit(
+        np.log(wav_k[scl_min:scl_max]), np.log(S1a[1, scl_min:scl_max]), deg=1, cov=True
+    )
+    yfitcw = np.exp(coefcw[1]) * wav_k[scl_min:scl_max] ** coefcw[0]
+    plt.plot(wav_k[scl_min:scl_max], yfitcw, "-", color="blue", linewidth=4, alpha=0.5)
+    print("Coherent Power law", coefcw[0])
+
+    # Fourier power law fit
+    limit = np.where((tab_k >= wav_k[fit_min]) & (tab_k <= wav_k[fit_max]))
+    coef, cov = np.polyfit(np.log(tab_k[limit]), np.log(spec_k[limit]), deg=1, cov=True)
+    yfit = np.exp(coef[1]) * tab_k[limit] ** coef[0]
+    plt.plot(tab_k[limit], yfit, "-", color="black", linewidth=2, alpha=0.5)
+    print("Fourier Power law", coef[0])
+
+    # show resolution limits
+    # Gaussian part not accurate below levelmin due to the way AMR works
+    sim_res_eff = 2 ** lvlmin / 10
+    sim_res_lvl_min = 2 ** lvlmin
+    plt.plot(
+        [sim_res_lvl_min / nsize, sim_res_lvl_min / nsize],
+        [1e-6, 1e8],
+        color="green",
+        ls="--",
+        label="dx(levelmin)",
+    )
+    plt.plot(
+        [sim_res_eff / nsize, sim_res_eff / nsize],
+        [1e-6, 1e8],
+        color="green",
+        ls="-",
+        label="dx(levelmin) x 10",
+    )
+
+    plt.xscale("log")
+    plt.yscale("log")
+    plt.xlabel(r"$k$ (pixel$^{-1}$)")
+    plt.ylabel(r"$P(k)$")
+    plt.ylim(1e-6, 1e8)
+    plt.legend(loc="lower left")
+    plt.savefig("powerspectrum_{}.png".format(label), bbox_inches="tight")
+    plt.close()
+
+
+def save_results(wt, S11a, wav_k, S1a, q, label):
+    np.save("wav_k_{}.npy".format(label), wav_k)
+    np.save("S1a_{}.npy".format(label), S1a)
+    np.save("wt_{}.npy".format(label), wt)
+    np.save("S11a_{}.npy".format(label), S11a)
+    np.save("q_{}.npy".format(label), q)
+
+
+def load_results(label):
+    wav_k = np.load("wav_k_{}.npy".format(label))
+    S1a = np.load("S1a_{}.npy".format(label))
+    wt = np.load("wt_{}.npy".format(label))
+    S11a = np.load("S11a_{}.npy".format(label))
+    q = np.load("q_{}.npy".format(label))
+    return wav_k, S1a, wt, S11a, q
+
+
+def analyse_sim(im):
+    """ Do the MnGseg analysis """
+    meanim = np.mean(im)
+    imzm = im - meanim
+    M = nb_scale(im.shape)
+
+    fit_min = 5  # minimal scale for fitting the power spectrum slope
+    fit_max = 11  # max scale
+    label = "final"  # label to identify parameter setup
+
+    # after a lot of trials I found a fixed q=2 is a good value
+    q = [2.0] * nb_scale(imzm.shape)
+    # wt, S11a, wav_k, S1a, q = fan_trans(imzm, reso=1, q=q, qdyn=False)
+    # alternatively you can let pywavan determine it automatically by setting skewl
+    # q=[3.0]*nb_scale(imzm.shape)
+    # wt, S11a, wav_k, S1a, q = fan_trans(imzm, reso=1, q=q, qdyn=True, skewl=0.4)
+    # print(q)
+
+    # save results because it can be long to calculate (especially if qdyn=True). Remark that wt and S11a are quite big
+    # save_results(wt, S11a, wav_k, S1a, q, label)
+    wav_k, S1a, wt, S11a, q = load_results(label)
+
+    # make images of the Gaussian and coherent part
+    make_images(im, wt, M, meanim, label)
+
+    # (optional) make the image for each scale
+    for s in range(fit_min, fit_max + 1):
+        scale_images(wt, M, label, scale=s)
+
+    # (optional) plot the histogram for each scale
+    plot_each_scale(S11a, wav_k, q, label, coherent=True)
+
+    # calc Fourier power spectrum for comparison
+    tab_k, spec_k = powspec(imzm, reso=1)
+
+    # plot the powerspectrum for each component and fit the slope
+    plot_components_power_spectrum(tab_k, wav_k, spec_k, S1a, label, fit_min, fit_max)