From a493839f4bd9f786ac67cd8879083600d0d43d0d Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Fri, 25 Jun 2021 09:21:41 +0200 Subject: [PATCH] [mngseg] add gaussian / coherent PS segementation tool --- mngseg.py | 275 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 275 insertions(+) create mode 100644 mngseg.py diff --git a/mngseg.py b/mngseg.py new file mode 100644 index 0000000..325c317 --- /dev/null +++ b/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)