Paper figures

Generated from AGORA simulations and trained DDPM checkpoint. Run on the cluster where the data files are available.

§8 plot-style rewrite (Wong palette, cividis cmaps, PDF output) is tracked separately.

import numpy as np
import healpy as hp
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import matplotlib.patches as patches
from matplotlib.lines import Line2D
from matplotlib.gridspec import GridSpec
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from scipy.ndimage import gaussian_filter1d
from pathlib import Path
from random import sample

from foregrounds_diffusion.flatmaps import map2cl, cl2map, get_lxly
from foregrounds_diffusion.plot_style import apply as _apply_style, WONG
from foregrounds_diffusion.preprocessing import (
    apply_stdnorm,
    renormalize_dm_maps,
    denormalize_dm_maps,
    load_all_moments,
)
from foregrounds_diffusion.statistics import stats

plt.rcParams.update({
    "text.usetex": True,
    "font.family": "serif",
    "font.serif": ["Computer Modern Roman", "DejaVu Serif"],
    "font.size": 12,
    "figure.dpi": 150,
    "savefig.dpi": 300,
    "savefig.bbox": "tight",
})

PROJECT_ROOT = Path("/home/apb86/cmb_foregrounds_diffusion")
DATA_DIR     = PROJECT_ROOT / "data"
PATCHES_DIR  = DATA_DIR / "low_pass" / "2mJy"
FIGURES_DIR  = PROJECT_ROOT / "plots" / "paper"
FIGURES_DIR.mkdir(parents=True, exist_ok=True)

PTSRC = 2
RES   = 256
flatskymapparams = [RES, RES, 1.40625, 1.40625]  # [nx, ny, dx_arcmin, dy_arcmin]

WONG = _apply_style(fig_width_pt=510.0, n_cols=2)

Fig 1

cib_150_map = hp.read_map(DATA_DIR / "mask_radio_cib_2mJy" / "mdpl2_150GHz_fullsky.fits")

hp.mollview(cib_150_map, title="", min=0, max=60, unit=r"$\mu K$", cmap='cividis')
fig = plt.gcf()
cax = fig.axes[-1]
cax.tick_params(labelsize=20)
plt.savefig("figures/cib_150_map.png", bbox_inches="tight")
plt.show()

cib_150_map_filtered = hp.read_map(DATA_DIR / "mask_radio_cib_2mJy" / "mdpl2_150GHz_fullsky_lmax7000.fits")

hp.mollview(cib_150_map_filtered, title="", min=0, max=60, unit=r"$\mu K$", cmap='cividis')
fig = plt.gcf()
cax = fig.axes[-1]
cax.tick_params(labelsize=20,)
plt.savefig("figures/cib_150_map_processed.png", bbox_inches="tight")
plt.show()

fname = PATCHES_DIR / "CIB_map_150GHz_256_st6_minmax_2mJy_zero_lp.npy"
processed_maps = np.load(fname);

fig, ax = plt.subplots(figsize=(5, 5))
im = ax.imshow(processed_maps[0], cmap='cividis', vmin=0, vmax=1)
ax.set_xlabel(r"$6^\circ$", fontsize=20)
ax.set_ylabel(r"$6^\circ$", fontsize=20)
ax.grid()

cbar = fig.colorbar(im, ax=ax, orientation='horizontal',
                    fraction=0.046)
cbar.set_ticks([0, 1])
cbar.ax.tick_params(labelsize=20)
ax.tick_params(axis='both', which='both',
               bottom=False, top=False,
               left=False, right=False,
               labelbottom=False, labelleft=False)

plt.savefig("figures/cib_150_map_flatsky.png", bbox_inches="tight")
plt.show()

fig, ax = plt.subplots(figsize=(12, 4))
ax.axis('off')

img_fullsky1 = mpimg.imread('figures/cib_150_map.png')
img_fullsky2 = mpimg.imread('figures/cib_150_map_processed.png')
img_flatsky  = mpimg.imread('figures/cib_150_map_flatsky.png')

nodes = {
    'node1': (0.05, 0.35, 0.1, 0.3),
    'node2': (0.43, 0.35, 0.1, 0.3),
    'node3': (0.67, 0.35, 0.25, 0.3),
}

def place_image(img, coords, zoom=0.1):
    x = coords[0] + coords[2] / 2
    y = coords[1] + coords[3] / 2
    im = OffsetImage(img, zoom=zoom)
    ab = AnnotationBbox(im, (x, y), frameon=False)
    ax.add_artist(ab)

place_image(img_fullsky1, nodes['node1'], zoom=0.1)
place_image(img_fullsky2, nodes['node2'], zoom=0.1)
place_image(img_flatsky,  nodes['node3'], zoom=0.1)

node1_bottom = (nodes['node1'][0] + nodes['node1'][2] / 2, nodes['node1'][1])
node2_bottom = (nodes['node2'][0] + nodes['node2'][2] / 2, nodes['node2'][1])
node3_bottom = (nodes['node3'][0] + nodes['node3'][2] / 2, nodes['node3'][1])

arrow_offset = -0.3
arrow_start12 = (node1_bottom[0] +0.1 , node1_bottom[1] + arrow_offset)
arrow_end12   = (node2_bottom[0]-0.1, node2_bottom[1] + arrow_offset)
arrow_start23 = (node2_bottom[0]+0.05, node2_bottom[1] + arrow_offset)
arrow_end23   = (node3_bottom[0]-0.05, node3_bottom[1] + arrow_offset)

ax.annotate("", xy=arrow_end12, xytext=arrow_start12,
            arrowprops=dict(arrowstyle="->", lw=1, shrinkA=5, shrinkB=5))
midpoint12 = ((arrow_start12[0] + arrow_end12[0]) / 2, (arrow_start12[1] + arrow_end12[1]) / 2)
ax.text(midpoint12[0], midpoint12[1] + 0.05, "2mJy point source mask", ha='center', fontsize=12)
ax.text(midpoint12[0], midpoint12[1] - 0.09, "Low pass filter at 7000", ha='center', fontsize=12)

ax.annotate("", xy=arrow_end23, xytext=arrow_start23,
            arrowprops=dict(arrowstyle="->", lw=1, shrinkA=5, shrinkB=5))
midpoint23 = ((arrow_start23[0] + arrow_end23[0]) / 2, (arrow_start23[1] + arrow_end23[1]) / 2)
ax.text(midpoint23[0], midpoint23[1] + 0.05, "Extracting flatsky patches", ha='center', fontsize=12)
ax.text(midpoint23[0], midpoint23[1] - 0.09, "Normalization 0 to 1", ha='center', fontsize=12)
plt.savefig("figures/map_processing.png", bbox_inches="tight")
plt.savefig("figures/map_processing.pdf", bbox_inches="tight")
plt.show()

Multifrequency maps

channels = ["CIB95", "CIB150", "CIB857"]
filenames = [PATCHES_DIR / f"cut_maps_RES_256_ANG_X_6.0_deg_2mJy_lp_{ch}.npy" for ch in channels]

train_maps_list = [np.load(fname) for fname in filenames]
train_maps = np.concatenate(train_maps_list, axis=-1)

dm_maps = np.load(PATCHES_DIR / "new_samples_14_cib_2mJy_zero_6x6_w_au_lp_three.npy")
dm_maps = renormalize_dm_maps(dm_maps, train_maps, variance_scaling=True)

example_idx = 100
example_agora = apply_stdnorm(train_maps[example_idx])
example_diffusion = apply_stdnorm(dm_maps[example_idx])
cmap='cividis'

plt.clf()
plt.figure(figsize=(10, 7))
plt.subplots_adjust(wspace=0.01)

titles = ["CIB 95 GHz", "CIB 150 GHz", "CIB 857 GHz"]
vmin_vmax_dic = {0: (-1, 4), 1: (-1, 4), 2: (-1, 4)}
#vmin_vmax_dic = {0: (None, None), 1: (None, None), 2: (None, None)}

sp_i = 1
for data, model in zip([example_agora, example_diffusion], ["Agora", "DDPM"]):
    for i in range(3):
        plt.subplot(2, 3, sp_i)
        vmin, vmax = vmin_vmax_dic[i]
        plt.imshow(data[:, :, i], cmap=cmap, aspect='auto', vmin=vmin, vmax=vmax)
        title(f"{model}: {titles[i]}", fontsize=20)
        axis('off')
        sp_i += 1

plt.tight_layout()
plt.savefig("figures/example_cib_triplet.pdf", bbox_inches="tight")
plt.savefig("figures/example_cib_triplet.png", bbox_inches="tight")
plt.show()

def correlation_coeff(cross, auto1, auto2):
    return np.array(cross) / np.sqrt(np.array(auto1) * np.array(auto2))

n_samples = 200
idxs = sample(range(dm_maps.shape[0]), n_samples)
cross_pairs = [(0, 1), (0, 2), (1, 2)]

auto_cls = {'train': {i: [] for i in range(3)},
            'dm':    {i: [] for i in range(3)}}
cross_cls = {'train': {p: [] for p in cross_pairs},
             'dm':    {p: [] for p in cross_pairs}}

ell, _ = map2cl(flatskymapparams, train_maps[0, :, :, 0])

for i in idxs:
    for tag, maps in zip(['train', 'dm'], [train_maps, dm_maps]):
        for ch in range(3):
            m = maps[i, :, :, ch]
            _, cl = map2cl(flatskymapparams, m)
            auto_cls[tag][ch].append(cl)
        for ch1, ch2 in cross_pairs:
            m1, m2 = maps[i, :, :, ch1], maps[i, :, :, ch2]
            _, cl = map2cl(flatskymapparams, m1, m2)
            cross_cls[tag][(ch1, ch2)].append(cl)

mean_auto_cls = {tag: {k: mean(v, axis=0) for k, v in auto.items()} for tag, auto in auto_cls.items()}
mean_cross_cls = {tag: {k: mean(v, axis=0) for k, v in cross.items()} for tag, cross in cross_cls.items()}

corr = {}
for tag in ['train', 'dm']:
    corr[tag] = {
        (ch1, ch2): mean(
            correlation_coeff(cross_cls[tag][(ch1, ch2)],
                              auto_cls[tag][ch1],
                              auto_cls[tag][ch2]), axis=0)
        for (ch1, ch2) in cross_pairs
    }

from matplotlib import colormaps as cm
pairs = [(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]
labels = ["95x95", "95x150", "95x857", "150x150", "150x857", "857x857"]
cmap = cm["RdBu_r"]
cmap_indices = [0.0, 0.1, 0.3, 0.7, 0.9, 1.0]
colors = [cmap(i) for i in cmap_indices]

fsval=14

plt.clf()
plt.figure(figsize=(10., 4.2))

plt.subplot(1, 2, 1)
for i, (ch1, ch2) in enumerate(pairs):
    if ch1 == ch2:
        cl_train = mean_auto_cls["train"][ch1]
        cl_dm = mean_auto_cls["dm"][ch1]
    else:
        cl_train = mean_cross_cls["train"][(ch1, ch2)]
        cl_dm = mean_cross_cls["dm"][(ch1, ch2)]

    plot(ell, (cl_train / cl_dm) - 1, lw=1, label=labels[i], color=colors[i])

axhline(0, color='gray', ls='--', lw=1)
xlabel(r"Multipole $\ell$")
plt.ylabel(r"$C_\ell^{\mathrm{Agora}} / C_\ell^{\mathrm{DDPM}} - 1$")
xlim(300, 4000)
ylim(-0.1, 0.1)
legend(fontsize=fsval, ncol=2)

plt.subplot(1, 2, 2)
for i, (ch1, ch2) in enumerate(pairs):
    if ch1 == ch2:
        continue

    plot(ell, corr["train"][(ch1, ch2)], lw=2, label=f"{labels[i]} (Agora)", color=colors[i])
    plot(ell, corr["dm"][(ch1, ch2)], lw=2, ls='--', label=f"{labels[i]} (DM)", color=colors[i])

xlabel(r"Multipole $\ell$")
plt.ylabel("Correlation Coefficient")
xlim(300, 4000)
ylim(0.65, 1.05)
axhline(1, color='gray', ls='--', lw=1)

handles = [
    Line2D([0], [0], color='k', lw=2, linestyle='-', label='Agora'),
    Line2D([0], [0], color='k', lw=2, linestyle='--', label='DDPM')
]

legend(
    handles=handles,
    #ncols=3,
    fontsize=fsval,
    #loc='upper right',
)

plt.tight_layout()
savefig("figures/cib_triplet_errors_and_correlation.pdf", bbox_inches="tight")
savefig("figures/cib_triplet_errors_and_correlation.png", bbox_inches="tight")
plt.show()

tSZ-CIB maps

filenames = [PATCHES_DIR / f"cut_maps_RES_256_ANG_X_6.0_deg_2mJy_lp_CIB150.npy",
            PATCHES_DIR / f"cut_maps_RES_256_ANG_X_6.0_deg_2mJy_lp_tsz3.npy"]

agora_maps_list = [np.load(fname) for fname in filenames]
agora_maps = np.concatenate(agora_maps_list, axis=-1)

num_samples = len(agora_maps)
num_train = int(0.8 * num_samples)
rng = np.random.default_rng(seed=42)
indices = rng.permutation(num_samples)
train_indices = indices[:num_train]
test_indices = indices[num_train:]
train_maps = agora_maps[train_indices]
test_maps = agora_maps[test_indices]

dm_maps = np.load(PATCHES_DIR / "new_samples_16_cib_tsz3_2mJy_zero_norm_6x6_w_au_lp.npy")
dm_maps_unscaled = renormalize_dm_maps(dm_maps, train_maps, variance_scaling=False)

dm_maps = np.load(PATCHES_DIR / "new_samples_16_cib_tsz3_2mJy_zero_norm_6x6_w_au_lp.npy")
dm_maps = renormalize_dm_maps(dm_maps, train_maps, variance_scaling=True)

#np.save("data/low_pass/2mJy/new_samples_16_cib_tsz3_2mJy_zero_norm_6x6_w_au_lp_varscaled_no_pickle.npy", dm_maps, allow_pickle=False)

filenames = [PATCHES_DIR / f"cut_maps_RES_256_ANG_X_6.0_deg_2mJy_lp_gaussian_cib_joint3.npy",
             PATCHES_DIR / f"cut_maps_RES_256_ANG_X_6.0_deg_2mJy_lp_gaussian_tsz_joint3.npy"]

gaussian_maps_list = [np.load(fname) for fname in filenames]
gaussian_maps = np.concatenate(gaussian_maps_list, axis=-1)

train_maps_std = apply_stdnorm(train_maps)
dm_maps_std = apply_stdnorm(dm_maps)

stats(train_maps[:,:,:,1])

example_idxs = [68, 132, 85, 100]

vmin_cib, vmax_cib = -1, 4
vmin_tsz, vmax_tsz = -10, 3
cmap_cib = 'cividis'
cmap_tsz = 'RdBu_r'
fsval = 16  # slightly larger y-axis labels
sim_labels = ["Sim1", "Sim2", "Sim3", "Sim4"]

plt.clf()
plt.figure(figsize=(10, 10))

for i, idx in enumerate(example_idxs):
    for row, data, cmap, vmin, vmax, label in zip(
        range(4),
        [train_maps_std, dm_maps_std, train_maps_std, dm_maps_std],
        [cmap_cib, cmap_cib, cmap_tsz, cmap_tsz],
        [vmin_cib, vmin_cib, vmin_tsz, vmin_tsz],
        [vmax_cib, vmax_cib, vmax_tsz, vmax_tsz],
        ["CIB Agora", "CIB DDPM", "tSZ Agora", "tSZ DDPM"]
    ):
        ax = plt.subplot(4, 4, i + 1 + row * 4)
        plt.imshow(data[idx, :, :, row // 2], cmap=cmap, vmin=vmin, vmax=vmax, aspect='auto')
        ax.set_xticks([])
        ax.set_yticks([])

        if i == 0:
            plt.ylabel(label, fontsize=fsval+2)


        ax.text(
            0.95, 0.85, sim_labels[i],
            transform=ax.transAxes,
            fontsize=14,
            color='black',
            ha='right',
            va='bottom',
            bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=1)
        )

plt.tight_layout()
savefig("figures/examples_grid_labeled.pdf", bbox_inches="tight")
savefig("figures/examples_grid_labeled.png", bbox_inches="tight")
plt.show()

example_idxs = [68, 132, 85, 100]

vmin_cib, vmax_cib = -1, 4
vmin_tsz, vmax_tsz = -10, 3
cmap_cib = 'cividis'
cmap_tsz = 'RdBu_r'
fsval = 16
sim_labels = ["Sim1", "Sim2", "Sim3", "Sim4"]

plt.clf()
plt.figure(figsize=(10, 10))

for i, idx in enumerate(example_idxs):
    for row, data, cmap, vmin, vmax, label in zip(
        range(4),
        [train_maps_std, dm_maps_std, train_maps_std, dm_maps_std],
        [cmap_cib, cmap_cib, cmap_tsz, cmap_tsz],
        [vmin_cib, vmin_cib, vmin_tsz, vmin_tsz],
        [vmax_cib, vmax_cib, vmax_tsz, vmax_tsz],
        ["CIB Agora", "CIB DDPM", "tSZ Agora", "tSZ DDPM"]
    ):
        ax = plt.subplot(4, 4, i + 1 + row * 4)
        plt.imshow(data[idx, :, :, row // 2], cmap=cmap, vmin=vmin, vmax=vmax, aspect='auto')
        ax.set_xticks([])
        ax.set_yticks([])

        if i == 0:
            plt.ylabel(label, fontsize=fsval + 2)

        ax.text(
            0.95, 0.85, sim_labels[i],
            transform=ax.transAxes,
            fontsize=14,
            color='black',
            ha='right',
            va='bottom',
            bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=1)
        )

# Add zoomed-in insets to subplot (3,2) and (3,3), which are subplot indices 10 and 11 (0-based)
for subplot_idx, zoom_coords in zip([14, 15], [(0.6, 0.6), (0.3, 0.3)]):
    ax_main = plt.gcf().axes[subplot_idx]

    # Create inset axes
    ax_inset = inset_axes(
    ax_main,
    width="30%", height="30%", loc='upper left',
    bbox_to_anchor=(1.05, 0.5, 0.3, 0.3),  # (x0, y0, width, height) in axes fraction
    bbox_transform=ax_main.transAxes,
    borderpad=0
)

    # Get image data and zoom in arbitrarily (adjust indices as needed)
    im_data = ax_main.images[0].get_array()
    zoom_slice = im_data[35:55, 35:55]  # You can change this

    ax_inset.imshow(
        zoom_slice,
        cmap=ax_main.images[0].get_cmap(),
        vmin=ax_main.images[0].get_clim()[0],
        vmax=ax_main.images[0].get_clim()[1],
        aspect='auto'
    )
    ax_inset.set_xticks([])
    ax_inset.set_yticks([])
    ax_inset.set_title("Zoom", fontsize=10)

    # Draw arrow from main plot to inset
    con = ConnectionPatch(
        xyA=zoom_coords, coordsA=ax_main.transAxes,
        xyB=(0.5, 0.5), coordsB=ax_inset.transAxes,
        axesA=ax_main, axesB=ax_inset,
        arrowstyle="->", color="white", linewidth=1.5
    )
    ax_main.add_artist(con)

plt.tight_layout()
plt.savefig("figures/examples_grid_labeled.pdf", bbox_inches="tight")
plt.savefig("figures/examples_grid_labeled.png", bbox_inches="tight")
plt.show()

fname = DATA_DIR / "ilc" / "ilc_weights_residuals_agora_fg_model.npy"
ilc_dict = np.load(fname, allow_pickle = True).item()
total_ilc_residuals_dict = ilc_dict['total_ilc_residuals']

l, nl_s4_wide = total_ilc_residuals_dict['s4_wide']['mv']
l, nl_s4_deep = total_ilc_residuals_dict['s4_deep']['mv']
l, nl_spt3g = total_ilc_residuals_dict['spt3g']['mv']

n_samples = 200

cls_test_cib_clean = []
cls_test_tsz_clean = []
cls_dm_cib_clean = []
cls_dm_tsz_clean = []
cls_dmu_cib_clean = []
cls_dmu_tsz_clean = []

cls_test_cib_noisy = []
cls_test_tsz_noisy = []
cls_dm_cib_noisy = []
cls_dm_tsz_noisy = []
cls_dmu_cib_noisy = []
cls_dmu_tsz_noisy = []

cls_test_cross_clean = []
cls_test_cross_noisy = []
cls_dm_cross_clean = []
cls_dm_cross_noisy = []
cls_dmu_cross_clean = []
cls_dmu_cross_noisy = []

for i in range(n_samples):
    noise = cl2map(flatskymapparams, nl_spt3g, l)

    cls_test_cib_clean.append(map2cl(flatskymapparams, test_maps[i, :, :, 0], test_maps[i, :, :, 0])[1])
    cls_test_tsz_clean.append(map2cl(flatskymapparams, test_maps[i, :, :, 1], test_maps[i, :, :, 1])[1])
    cls_test_cross_clean.append(map2cl(flatskymapparams, test_maps[i, :, :, 0], test_maps[i, :, :, 1])[1])

    noisy_cib = test_maps[i, :, :, 0] + noise
    noisy_tsz = test_maps[i, :, :, 1] + noise

    cls_test_cib_noisy.append(map2cl(flatskymapparams, noisy_cib, noisy_cib)[1])
    cls_test_tsz_noisy.append(map2cl(flatskymapparams, noisy_tsz, noisy_tsz)[1])
    cls_test_cross_noisy.append(map2cl(flatskymapparams, noisy_cib, noisy_tsz)[1])

    ###
    cls_dm_cib_clean.append(map2cl(flatskymapparams, dm_maps[i, :, :, 0], dm_maps[i, :, :, 0])[1])
    cls_dm_tsz_clean.append(map2cl(flatskymapparams, dm_maps[i, :, :, 1], dm_maps[i, :, :, 1])[1])
    cls_dm_cross_clean.append(map2cl(flatskymapparams, dm_maps[i, :, :, 0], dm_maps[i, :, :, 1])[1])

    noisy_dm_cib = dm_maps[i, :, :, 0] + noise
    noisy_dm_tsz = dm_maps[i, :, :, 1] + noise

    cls_dm_cib_noisy.append(map2cl(flatskymapparams, noisy_dm_cib, noisy_dm_cib)[1])
    cls_dm_tsz_noisy.append(map2cl(flatskymapparams, noisy_dm_tsz, noisy_dm_tsz)[1])
    cls_dm_cross_noisy.append(map2cl(flatskymapparams, noisy_dm_cib, noisy_dm_tsz)[1])

    ###
    cls_dmu_cib_clean.append(map2cl(flatskymapparams, dm_maps_unscaled[i, :, :, 0], dm_maps_unscaled[i, :, :, 0])[1])
    cls_dmu_tsz_clean.append(map2cl(flatskymapparams, dm_maps_unscaled[i, :, :, 1], dm_maps_unscaled[i, :, :, 1])[1])
    cls_dmu_cross_clean.append(map2cl(flatskymapparams, dm_maps_unscaled[i, :, :, 0], dm_maps_unscaled[i, :, :, 1])[1])

    noisy_dmu_cib = dm_maps_unscaled[i, :, :, 0] + noise
    noisy_dmu_tsz = dm_maps_unscaled[i, :, :, 1] + noise

    cls_dmu_cib_noisy.append(map2cl(flatskymapparams, noisy_dmu_cib, noisy_dmu_cib)[1])
    cls_dmu_tsz_noisy.append(map2cl(flatskymapparams, noisy_dmu_tsz, noisy_dmu_tsz)[1])
    cls_dmu_cross_noisy.append(map2cl(flatskymapparams, noisy_dmu_cib, noisy_dmu_tsz)[1])

el, _ = map2cl(flatskymapparams, test_maps[0, :, :, 0])
dl_factor = el * (el + 1) / (2 * np.pi)

mean_test_cib  = dl_factor * np.mean(cls_test_cib_clean, axis=0)
mean_test_tsz  = dl_factor * np.mean(cls_test_tsz_clean, axis=0)
mean_test_cross = dl_factor * np.mean(cls_test_cross_clean, axis=0)

mean_dm_cib    = dl_factor * np.mean(cls_dm_cib_clean, axis=0)
mean_dm_tsz    = dl_factor * np.mean(cls_dm_tsz_clean, axis=0)
mean_dm_cross   = dl_factor * np.mean(cls_dm_cross_clean, axis=0)

mean_dmu_cib    = dl_factor * np.mean(cls_dmu_cib_clean, axis=0)
mean_dmu_tsz    = dl_factor * np.mean(cls_dmu_tsz_clean, axis=0)
mean_dmu_cross   = dl_factor * np.mean(cls_dmu_cross_clean, axis=0)

std_test_cib   = dl_factor * np.std(cls_test_cib_noisy, axis=0)
std_test_tsz   = dl_factor * np.std(cls_test_tsz_noisy, axis=0)
std_test_cross = dl_factor * np.std(cls_test_cross_noisy, axis=0)

bias_cross = (mean_test_cross - mean_dm_cross) / std_test_cross
bias_crossu = (mean_test_cross - mean_dmu_cross) / std_test_cross

plt.clf()
plt.figure(figsize=(10, 5))
gs = GridSpec(2, 2, height_ratios=[3, 1], hspace=0.0)

colors = {'CIB': WONG[5], 'tSZ': WONG[6], 'cross': WONG[7]}
fsval = 14

ax1 = plt.subplot(gs[0, 0])
fill_between(el, mean_test_cib - std_test_cib, mean_test_cib + std_test_cib, color=colors['CIB'], alpha=0.2)
plot(el, mean_test_cib, color=colors['CIB'], lw=1, ls='-', label="CIB (Agora)")
plot(el, mean_dm_cib, color=colors['CIB'], lw=1, ls='--', label="CIB (DDPM)")
#plot(el, mean_dmu_cib, color=colors['CIB'], lw=1, ls=':', )


fill_between(el, mean_test_tsz - std_test_tsz, mean_test_tsz + std_test_tsz, color=colors['tSZ'], alpha=0.2)
plot(el, mean_test_tsz, color=colors['tSZ'], lw=1, ls='-', label="tSZ (Agora)")
plot(el, mean_dm_tsz, color=colors['tSZ'], lw=1,ls='--', label="tSZ (DDPM)")
#plot(el, mean_dmu_tsz, color=colors['tSZ'], lw=1,ls=':')


xlim(300, 4200)
ylim(-1, 30)
plt.ylabel(r"$\ell(\ell+1)/2\pi\ C_\ell \ \left[\mu K^2\right]$", fontsize=fsval)
tick_params(axis='both', labelsize=fsval)
legend(fontsize=fsval, loc='upper left', frameon=True)

ax2 = plt.subplot(gs[1, 0], sharex=ax1)
bias_cib = (mean_test_cib - mean_dm_cib) / std_test_cib
bias_tsz = (mean_test_tsz - mean_dm_tsz) / std_test_tsz
plot(el, bias_cib, color=colors['CIB'], lw=1)
plot(el, bias_tsz, color=colors['tSZ'], lw=1)

axhline(0, color='gray', ls='--', lw=1)
xlabel(r"Multipole $\ell$", fontsize=fsval)
plt.ylabel(r"$\Delta C_\ell / \sigma C_\ell$", fontsize=fsval)
ylim(-0.4, 0.4)
tick_params(axis='both', labelsize=fsval)


ax3 = plt.subplot(gs[0, 1])
fill_between(el, mean_test_cross - std_test_cross, mean_test_cross + std_test_cross,
             color=colors['cross'], alpha=0.2)

plot(el, mean_test_cross, color=colors['cross'], lw=1, ls='-', label="CIB × tSZ (Agora)")
plot(el, mean_dm_cross, color=colors['cross'], lw=1, ls='--', label="CIB × tSZ (DDPM)")
#plot(el, mean_dmu_cross, color=colors['cross'], lw=1, ls=':',)

xlim(300, 4200)
ylim(-2.9, 4)
#plt.ylabel(r"$\ell(\ell+1)/2\pi\ C_\ell$", fontsize=fsval)
tick_params(axis='both', labelsize=fsval)
legend(fontsize=fsval, loc='upper left', frameon=True)

ax4 = plt.subplot(gs[1, 1], sharex=ax3)
plot(el, bias_cross, color=colors['cross'], lw=1)
axhline(0, color='gray', ls='--', lw=1)
xlabel(r"Multipole $\ell$", fontsize=fsval)
ylim(-0.4, 0.4)
tick_params(axis='both', labelsize=fsval)

plt.tight_layout()
plt.savefig("figures/power_spectra_comparison_cib_tsz_combined.pdf", bbox_inches="tight")
plt.savefig("figures/power_spectra_comparison_cib_tsz_combined.png", bbox_inches="tight")
plt.show()

def smooth_hist(pixels, bins):
    hist, _ = np.histogram(pixels, bins=bins, density=True)
    return gaussian_filter1d(hist, sigma=2)

pixels_training_cib = train_maps[:n_samples, :, :, 0].flatten()
pixels_training_tsz = -1*train_maps[:n_samples, :, :, 1].flatten()
pixels_samples_cib  = dm_maps[:n_samples, :, :, 0].flatten()
pixels_samples_tsz  = -1*dm_maps[:n_samples, :, :, 1].flatten()
pixels_samplesu_cib  = dm_maps_unscaled[:n_samples, :, :, 0].flatten()
pixels_samplesu_tsz  = -1*dm_maps_unscaled[:n_samples, :, :, 1].flatten()

#bins_cib = logspace(0., 2, 500)
#bins_tsz = logspace(0., 2, 500)
bins_cib = linspace(0., 70, 500)
bins_tsz = linspace(0., 90, 500)

bin_centers_cib = 0.5 * (bins_cib[:-1] + bins_cib[1:])
train_hist_cib  = smooth_hist(pixels_training_cib, bins_cib)
gen_hist_cib    = smooth_hist(pixels_samples_cib, bins_cib)
genu_hist_cib    = smooth_hist(pixels_samplesu_cib, bins_cib)

bin_centers_tsz = 0.5 * (bins_tsz[:-1] + bins_tsz[1:])
train_hist_tsz  = smooth_hist(pixels_training_tsz, bins_tsz)
gen_hist_tsz    = smooth_hist(pixels_samples_tsz, bins_tsz)
genu_hist_tsz    = smooth_hist(pixels_samplesu_tsz, bins_tsz)

plt.clf()
fig, axs = subplots(2, 2, figsize=(8, 6), sharex='col', sharey='row', gridspec_kw={'wspace': 0, 'hspace': 0})

fsval = 14
colors = {'DM': WONG[5], 'train': 'black'}

axs[0, 0].plot(bin_centers_cib, train_hist_cib, color=colors['train'], lw=1, label='Agora')
axs[0, 0].plot(bin_centers_cib, gen_hist_cib,   color=colors['DM'], lw=1, label='DDPM')
axs[0, 0].set_ylabel("Normalized counts", fontsize=fsval)
axs[0, 0].set_title("CIB one-point PDF", fontsize=fsval)
axs[0, 0].legend(fontsize=fsval)
axs[0, 0].tick_params(labelbottom=False)

axs[0, 1].plot(bin_centers_tsz, train_hist_tsz, color=colors['train'], lw=1, label='Agora')
axs[0, 1].plot(bin_centers_tsz, gen_hist_tsz,   color=colors['DM'], lw=1, label='DDPM')
axs[0, 1].set_title("tSZ one-point PDF", fontsize=fsval)
axs[0, 1].tick_params(labelbottom=False, labelleft=False)

axs[1, 0].plot(bin_centers_cib, train_hist_cib, color=colors['train'], lw=1)
axs[1, 0].plot(bin_centers_cib, gen_hist_cib,   color=colors['DM'], lw=1)
axs[1, 0].set_yscale('log')
axs[1, 0].set_xlabel("Pixel Intensity", fontsize=fsval)
axs[1, 0].set_ylabel("Normalized counts", fontsize=fsval)

axs[1, 1].plot(bin_centers_tsz, train_hist_tsz, color=colors['train'], lw=1)
axs[1, 1].plot(bin_centers_tsz, gen_hist_tsz,   color=colors['DM'], lw=1)
axs[1, 1].set_yscale('log')
axs[1, 1].set_xlabel("Pixel Intensity", fontsize=fsval)
axs[1, 1].tick_params(labelleft=False)


plt.tight_layout()
plt.savefig("figures/histogram_comparison_curve_2panel.pdf", bbox_inches="tight")
plt.savefig("figures/histogram_comparison_curve_2panel.png", bbox_inches="tight")
plt.show()

plt.clf()
fig, axs = subplots(1, 2, figsize=(8, 3), sharey=True, gridspec_kw={'wspace': 0})

fsval = 14
colors = {'DM': 'steelblue', 'train': 'black','Unscaled DM': 'red'}

axs[0].plot(-bin_centers_tsz, train_hist_tsz, color=colors['train'], lw=1.5, label='Agora')
axs[0].plot(-bin_centers_tsz, gen_hist_tsz,   color=colors['DM'], lw=1.5, label='DDPM')
axs[0].plot(-bin_centers_tsz, genu_hist_tsz,   color=colors['Unscaled DM'], lw=1, alpha=0.7, ls='--', label='Unscaled DDPM')
axs[0].set_yscale('log')
axs[0].set_xlabel(r"Pixel Intensity [$\mu K$]", fontsize=fsval)
axs[0].set_ylabel("Normalized counts", fontsize=fsval)
axs[0].set_ylim(5e-9,5e-1)
axs[0].set_xlim(-90,0)

axs[0].set_title("tSZ one-point PDF", fontsize=fsval)
axs[0].legend(fontsize=fsval)

axs[1].plot(bin_centers_cib, train_hist_cib, color=colors['train'], lw=1.5)
axs[1].plot(bin_centers_cib, gen_hist_cib,   color=colors['DM'], lw=1.5)
axs[1].plot(bin_centers_cib, genu_hist_cib,   color=colors['Unscaled DM'], lw=1, alpha=0.7, ls="--")
axs[1].set_yscale('log')
axs[1].set_ylim(1e-7,4e-1)
axs[1].set_xlim(0,70)

axs[1].set_xlabel(r"Pixel Intensity [$\mu K$]", fontsize=fsval)
axs[1].set_title("CIB one-point PDF", fontsize=fsval)
axs[1].tick_params(labelleft=False)

plt.tight_layout()
savefig("figures/histogram_comparison.pdf", bbox_inches="tight")
savefig("figures/histogram_comparison.png", bbox_inches="tight")
plt.show()

Moments

labels = ["Bispectrum $S_{3}$ [$\mu$K$^{3}$]",  "Trispectrum $S_{4}$ [$\mu$K$^{4}$]"]
title_labels = [
    r"$S_3$",r"$S_4$",
]


moment_keys = ["moment_03", "moment_07"]
mul_fac_dic = {0: 1e4, 1: 1e6}

noise_levels = ["s4deep_noise"]

models = ["train", "diffusion", "gaussian"]
model_labels = ["Agora", "DDPM", "Gaussian"]
markers = {"train": ".", "diffusion": ".", "gaussian": "."}

noise_colors = {
    #"spt3g_noise": "goldenrod",
    #"s4wide_noise": "steelblue",
    "s4deep_noise": "mediumvioletred"
}

colors = {
    "gaussian": "orangered",
    "train": "black",
    "diffusion": "steelblue"
}

lmin = 300
lmax = 6000
delta_l = 720
bp_edges = np.arange(lmin, lmax, delta_l)
bandpass_centers = (bp_edges[:-1] + bp_edges[1:]) // 2

def load_moment_sum(model, noise=None):
    tag = "samples_curve" if model == "diffusion" else model
    suffix = f"_sum_{noise}" if noise else "_sum"
    fname = PATCHES_DIR / f"moments_{tag}_2mJy_deltaell_720_200rlz_6x6_lp_joint3{suffix}.npy"
    return load_all_moments(fname, bandpass_centers)

sum_data_clean = {
    model: load_moment_sum(model)
    for model in models
}

sum_data_noisy = {
    noise: {
        model: load_moment_sum(model, noise)
        for model in models
    }
    for noise in noise_levels
}

plt.clf()

tr, tc = 1, 2
plt.figure(figsize=(8., 4.2))
plt.subplots_adjust(wspace=0.1)

alphaval = 1.
mewval = 0.8
fsval = 14
msval = 8.
capsizeval = 0.

offset_step = 0.
model_offset_step = 50
bpwidth = np.diff(bandpass_centers)[0]

for idx, (label, moment_key) in enumerate(zip(title_labels, moment_keys)):
    ax = plt.subplot(tr, tc, idx + 1)

    for noise_i, noise_level in enumerate(noise_levels):
        base_offset = (noise_i - len(noise_levels) // 2) * offset_step

        for model_i, model in enumerate(markers):
            color = colors[model]
            mean_data = np.array(sum_data_clean[model][moment_key])
            noisy_data = np.array(sum_data_noisy[noise_level][model][moment_key])

            mean_vals = mean(mean_data, axis=0)
            std_errs = std(noisy_data, axis=0) / np.sqrt(noisy_data.shape[0])
            offset = base_offset + (model_i - len(models) // 2) * model_offset_step
            shifted_x = bandpass_centers[1:] + offset

            errorbar(shifted_x, mean_vals[1:]* mul_fac_dic[idx], yerr=std_errs[1:]* mul_fac_dic[idx],
                     fmt=markers[model], capsize=capsizeval, color=color, alpha=alphaval,
                     label=model_labels[model_i], ms=msval, lw=mewval)

    if idx == 0:
        plt.ylabel(labels[0], fontsize=fsval)
        #ylim(-8e-4, 0.5e-4)
        axhline(0., lw=1., alpha=0.5)
    elif idx == 1:
        plt.ylabel(labels[1], fontsize=fsval)
        legend(fontsize=12, numpoints=1, handletextpad=-0.2)
        #ylim(-0.5e-6, 8e-6)
        axhline(0., lw=1., alpha=0.5)

    #title(label, fontsize=fsval)
    xlim(bandpass_centers[1] - bpwidth/2 - 20, max(bandpass_centers) + bpwidth/2 + 30)
    xlabel(r"Multipole $\ell_c$", fontsize=fsval)
    grid(True, which='both', ls='-', lw=0.2, alpha=0.2)

    for cntr in range(len(bandpass_centers)):
        if cntr % 2 != 0:
            x1 = bandpass_centers[cntr] - bpwidth / 2
            x2 = bandpass_centers[cntr] + bpwidth / 2
            axvspan(x1, x2, color='peru', alpha=0.2)

        label_with_mulfac = r'%s [$10^{%g}$]' %(label, np.log10(mul_fac_dic[idx]))
        ax.set_title(label_with_mulfac, fontsize=fsval)

    grid(False, which='both', axis = 'both')
plt.tight_layout()
savefig("/global/homes/k/kp22/ddpm_paper/figures/moments_sum_cib_tsz.pdf", bbox_inches="tight")
savefig("/global/homes/k/kp22/ddpm_paper/figures/moments_sum_cib_tsz.png", bbox_inches="tight")
plt.show()

noise_levels = ["spt3g_noise", "s4wide_noise", "s4deep_noise"]
models = ["gaussian", "train", "diffusion"]

base_path = PATCHES_DIR / ""
file_template_clean = base_path + "moments_{tag}_2mJy_deltaell_720_200rlz_6x6_lp_joint3.npy"
file_template_noisy = base_path + "moments_{tag}_2mJy_deltaell_720_200rlz_6x6_lp_joint3_{noise}.npy"

means_clean = {
    "train": load_all_moments(file_template_clean.format(tag="train"), bandpass_centers),
    "diffusion": load_all_moments(file_template_clean.format(tag="samples_curve"), bandpass_centers),
    "gaussian": load_all_moments(file_template_clean.format(tag="gaussian"), bandpass_centers),
}

errors_noisy_all = {
    noise: {
        "train": load_all_moments(file_template_noisy.format(tag="train", noise=noise), bandpass_centers),
        "diffusion": load_all_moments(file_template_noisy.format(tag="samples_curve", noise=noise), bandpass_centers),
        "gaussian": load_all_moments(file_template_noisy.format(tag="gaussian", noise=noise), bandpass_centers),
    }
    for noise in noise_levels
}

exclude_n = 1
fsval = 16
offset_step = 100

labels = [
    r"$S_3^{aaa}$", r"$S_3^{bbb}$", r"$S_3^{aab}$", r"$S_3^{abb}$",
    r"$S_4^{aaaa}$", r"$S_4^{bbbb}$", r"$S_4^{aaab}$", r"$S_4^{aabb}$", r"$S_4^{abbb}$"
]

colors = {"gaussian": "orangered", "train": "black", "diffusion": "steelblue"}
markers = {"s4deep_noise": "o", "s4wide_noise": "^", "spt3g_noise": "*"}

bpwidth = np.diff(bandpass_centers)[0]

plt.clf()
fig, axes = subplots(3, 3, figsize=(16, 12), sharex=True)
axes = axes.flatten()

for idx, label in enumerate(labels):
    ax = axes[idx]
    moment_key = f"moment_{3 + idx:02d}"

    for noise_i, noise_level in enumerate(noise_levels):
        offset = (noise_i - len(noise_levels) // 2) * offset_step

        for model in models:
            mean_vals = mean(means_clean[model][moment_key], axis=0)
            std_errs  = std(errors_noisy_all[noise_level][model][moment_key], axis=0) / np.sqrt(len(errors_noisy_all[noise_level][model][moment_key]))

            shifted_x = bandpass_centers[exclude_n:] + offset
            label_str = f"{model.capitalize()}" if idx == 0 and noise_i == 0 else None

            ax.errorbar(shifted_x, mean_vals[exclude_n:], yerr=std_errs[exclude_n:],
                        fmt=markers[noise_level], capsize=2, color=colors[model],
                        alpha=0.5, label=label_str)

    for cntr in range(len(bandpass_centers)):
        if cntr % 2 != 0:
            x1 = bandpass_centers[cntr] - bpwidth / 2
            x2 = bandpass_centers[cntr] + bpwidth / 2
            ax.axvspan(x1, x2, color='peru', alpha=0.2)

    ax.set_title(label, fontsize=fsval)
    ax.grid(True, which='both', ls='-', lw=0.2, alpha=0.2)
    if idx // 3 == 2:
        ax.set_xlabel(r"Multipole $\ell_c$", fontsize=fsval)
    if idx % 3 == 0:
        ax.set_ylabel("Statistic value", fontsize=fsval)

legend_elements = (
    [Line2D([0], [0], color=colors[m], lw=4, label=m.capitalize()) for m in models] +
    [Line2D([0], [0], marker=markers[n], color='black', linestyle='None', markersize=8, label=n.replace('_', '-')) for n in markers] +
    [Line2D([], [], color='none', label=r"$a$ = CIB,\ $b$ = tSZ")]
)
fig.legend(handles=legend_elements, fontsize=14, loc='lower center',
           bbox_to_anchor=(0.5, -0.05), ncol=len(legend_elements), frameon=False)

plt.tight_layout()
plt.savefig("figures/moments_joint_cib_tsz.pdf", bbox_inches="tight")
plt.savefig("figures/moments_joint_cib_tsz.png", bbox_inches="tight")
plt.show()

exclude_n = 1
fsval = 16
offset_step = 100

labels = [
    #r"$S_3^{aaa}$", r"$S_3^{aab}$", r"$S_3^{abb}$", r"$S_3^{bbb}$",
    r"$S_3^{aaa}$", r"$S_3^{bbb}$", r"$S_3^{aab}$", r"$S_3^{abb}$",
]

colors = {"gaussian": "orangered", "train": "black", "diffusion": "steelblue"}
markers = {"s4deep_noise": "o", "s4wide_noise": "^", "spt3g_noise": "*"}
expname_dic = {"s4deep_noise": r'S4-Ultra deep', "s4wide_noise": r"S4-Wide", "spt3g_noise": r"SPT-3G"}
model_labels = {"gaussian": "Gaussian", "train": "Agora", "diffusion": "DDPM"}

bpwidth = np.diff(bandpass_centers)[0]

tr, tc = 2, 2
plt.clf()
fig, axes = subplots(tr, tc, figsize=(8, 6.5), sharex=True)
axes = axes.flatten()
plt.subplots_adjust(wspace = 0.02)

mul_fac_dic = {0: 1e4, 1: 1e4, 2: 1e5, 3: 1e5}

for idx, label in enumerate(labels):
    ax = axes[idx]
    moment_key = f"moment_{3 + idx:02d}"
    print(moment_key)
    curr_mul_fac = mul_fac_dic[idx]

    for noise_i, noise_level in enumerate(noise_levels):
        offset = (noise_i - len(noise_levels) // 2) * offset_step

        for model in models:
            mean_vals = mean(means_clean[model][moment_key], axis=0)
            std_errs  = std(errors_noisy_all[noise_level][model][moment_key],axis=0) / np.sqrt(len(errors_noisy_all[noise_level][model][moment_key]))

            shifted_x = bandpass_centers[exclude_n:] + offset
            label_str = f"{model_labels[model]}" if idx == 0 and noise_i == 0 else None

            ax.errorbar(shifted_x, mean_vals[exclude_n:] * curr_mul_fac, yerr=std_errs[exclude_n:]* curr_mul_fac,
                        fmt=markers[noise_level], capsize=2, color=colors[model],
                        alpha=0.5,
                        #label=label_str,
                        ms = 4.)


    for cntr in range(len(bandpass_centers)):
        if cntr % 2 != 0:
            x1 = bandpass_centers[cntr] - bpwidth / 2
            x2 = bandpass_centers[cntr] + bpwidth / 2
            ax.axvspan(x1, x2, color='peru', alpha=0.2)

    ax.axhline(0., lw = 0.5, color = 'black')

    label_with_mulfac = r'%s [$10^{%g}$]' %(label, np.log10(curr_mul_fac))
    ax.set_title(label_with_mulfac, fontsize=fsval)
    ax.grid(False, which='both', axis = 'both')#ls='-', lw=0.2, alpha=0.2)
    if idx >= tc:
        ax.set_xlabel(r"Multipole $\ell_c$", fontsize=fsval)
    if idx % tc == 0:
        #ax.set_ylabel("Statistic value", fontsize=fsval)
        ax.set_ylabel("Bispectra $S_{3}$ [$\mu$K$^{3}$]", fontsize=fsval)

    if idx == 0: #add legend for models
        for model in models:
            print(model)
            label_str = f"{model_labels[model]}"
            ax.errorbar([], [], yerr=[0],
                        fmt='o',
                        capsize=2,
                        color=colors[model],
                        alpha=0.5,
                        label=label_str,
                        ms = 4.)
        ax.legend(loc = 1, fontsize = fsval - 3, numpoints = 1)
        ax.text(4000, 1.75, r'$a = {\rm CIB}$')
        ax.text(4000, 1.65, r'$b = {\rm tSZ}$')
    if idx == 1: #add legend for models
        for expname in expname_dic:
            label_str = r'%s' %(expname_dic[expname])
            ax.errorbar([], [], yerr=[0],
                        capsize=2,
                        color='black',
                        fmt=markers[expname],
                        alpha=0.5,
                        label=label_str,
                        ms = 4.,
                        )
        leg = ax.legend( fontsize = fsval - 3, title = r'{\bf Noise level:}',
                        title_fontsize = fsval - 3, numpoints = 1)
        leg._legend_box.align = "left"

plt.tight_layout()
savefig("figures/bispectra_joint_cib_tsz.pdf", bbox_inches="tight")
savefig("figures/bispectra_joint_cib_tsz.png", bbox_inches="tight")
plt.show()

exclude_n = 1
fsval = 13
offset_step = 100

subplot_panel_dic = {
    #value for each key is (rowno, colno), rowspan, colspan, mulfac.
    r"$S_4^{aaaa}$": [(0, 0), 3, 1, 1e7],
    r"$S_4^{bbbb}$": [(3, 0), 3, 1, 1e6],
    r"$S_4^{aaab}$": [(0, 1), 2, 1, 1e8],
    r"$S_4^{aabb}$": [(2, 1), 2, 1, 1e8],
    r"$S_4^{abbb}$": [(4, 1), 2, 1, 1e7],
}

colors = {"gaussian": "orangered", "train": "black", "diffusion": "steelblue"}
markers = {"s4deep_noise": "o", "s4wide_noise": "^", "spt3g_noise": "*"}
expname_dic = {"s4deep_noise": r'S4-Ultra deep', "s4wide_noise": r"S4-Wide", "spt3g_noise": r"SPT-3G"}

bpwidth = np.diff(bandpass_centers)[0]


plt.clf()
tr, tc = 6, 2
fig = plt.figure(figsize=(9.2, 9.2))
plt.subplots_adjust(wspace = 0.15, hspace = 0.4)

for idx,label in enumerate(subplot_panel_dic):
    curr_sbpl, curr_rowspan, curr_colspan, curr_mul_fac = subplot_panel_dic[label]
    curr_row, curr_col = curr_sbpl
    ax = subplot2grid( (tr, tc), curr_sbpl, colspan = curr_colspan, rowspan = curr_rowspan)

    moment_key = f"moment_{7 + idx:02d}"
    print(moment_key)
    for noise_i, noise_level in enumerate(noise_levels):
        offset = (noise_i - len(noise_levels) // 2) * offset_step

        for model in models:
            mean_vals = mean(means_clean[model][moment_key], axis=0)
            std_errs  = std(errors_noisy_all[noise_level][model][moment_key], axis=0) / np.sqrt(len(errors_noisy_all[noise_level][model][moment_key]))

            shifted_x = bandpass_centers[exclude_n:] + offset
            label_str = f"{model.capitalize()}" if idx == 0 and noise_i == 0 else None

            ax.errorbar(shifted_x, mean_vals[exclude_n:] * curr_mul_fac, yerr=std_errs[exclude_n:]* curr_mul_fac,
                        fmt=markers[noise_level], capsize=2, color=colors[model],
                        alpha=0.5,
                        #label=label_str,
                        ms = 4.)

    for cntr in range(len(bandpass_centers)):
        if cntr % 2 != 0:
            x1 = bandpass_centers[cntr] - bpwidth / 2
            x2 = bandpass_centers[cntr] + bpwidth / 2
            ax.axvspan(x1, x2, color='peru', alpha=0.2)

    ax.axhline(0., lw = 0.5, color = 'black')

    label_with_mulfac = r'%s [$10^{%g}$]' %(label, np.log10(curr_mul_fac))
    ax.set_title(label_with_mulfac, fontsize=fsval)
    ax.grid(False, which='both', axis = 'both')#ls='-', lw=0.2, alpha=0.2)
    if curr_row in [3, 4]:
        ax.set_xlabel(r"Multipole $\ell_c$", fontsize=fsval)
    else:
        setp(ax.get_xticklabels(), visible=False);
    if curr_col == 0:
        ax.set_ylabel("Trispectra $S_{4}$ [$\mu$K$^{4}$]", fontsize=fsval)

    for label in ax.get_xticklabels(): label.set_fontsize(fsval-2)
    for label in ax.get_yticklabels(): label.set_fontsize(fsval-2)

    if curr_sbpl in [(0,0)]: #add legend for models
        for model in models:
            print(model)
            label_str = f"{model_labels[model]}"
            ax.errorbar([], [], yerr=[0],
                        fmt='o',
                        capsize=2,
                        color=colors[model],
                        alpha=0.5,
                        label=label_str,
                        ms = 4.)
        ax.legend( fontsize = fsval - 2, numpoints = 1)
        ax.text(4000,2.65, r'$a = {\rm CIB}$')
        ax.text(4000,2.5, r'$b = {\rm tSZ}$')

    if curr_sbpl in [(3,0)]: #add legend for experiments
        for expname in expname_dic:
            label_str = r'%s' %(expname_dic[expname])
            ax.errorbar([], [], yerr=[0],
                        capsize=2,
                        color='black',
                        fmt=markers[expname],
                        alpha=0.5,
                        label=label_str,
                        ms = 4.,
                        )
        leg = ax.legend( fontsize = fsval - 2, title = r'{\bf Noise level:}',
                        title_fontsize = fsval - 2, numpoints = 1)
        leg._legend_box.align = "left"

savefig("figures/trispectra_joint_cib_tsz.pdf", bbox_inches="tight")
savefig("figures/trispectra_joint_cib_tsz.png", bbox_inches="tight")
plt.show()

Minkowski Functionals

loaded = np.load(PATCHES_DIR / "minkowski_results.npz", allow_pickle=True)
results = {k: loaded[k].tolist() for k in loaded}
thresholds = np.linspace(0., 1, 50)

fsval = 16
fig, ax = plt.subplots(3, 2, figsize=(10, 8))

titles = ["M0 (Area)", "M1 (Perimeter)", "M2 (Genus)"]
column_titles = ["CIB", "tSZ"]

for i, stat in enumerate(['M0', 'M1', 'M2']):
    for j, label in enumerate(['cib', 'tsz']):
        a = ax[i, j]
        train_data = results[f'{stat}_train_{label}']
        samples_data = results[f'{stat}_samples_{label}']
        gaussian_data = results[f'{stat}_gaussian_{label}']

        train_means = [x[0] for x in train_data]
        train_stds = [x[1] for x in train_data]
        samples_means = [x[0] for x in samples_data]
        samples_stds = [x[1] for x in samples_data]
        gaussian_means = [x[0] for x in gaussian_data]
        gaussian_stds = [x[1] for x in gaussian_data]

        #if label=='tsz':
        #a.set_xscale('log')
        #if label == 'tsz':
        #    a.set_xlim(0.8, 1)

        if stat == 'M0':
            a.set_yscale('log')

        a.plot(thresholds, train_means, label='Agora', color='black', lw=3, alpha=0.7)
        a.fill_between(thresholds, np.array(train_means) - np.array(train_stds),
                       np.array(train_means) + np.array(train_stds), alpha=0.3, color='black')
        a.plot(thresholds, samples_means, label='DDPM', color='steelblue', lw=1)
        a.fill_between(thresholds, np.array(samples_means) - np.array(samples_stds),
                       np.array(samples_means) + np.array(samples_stds), alpha=0.3, color='steelblue')
        a.plot(thresholds, gaussian_means, label='Gaussian', color=WONG[6], lw=1)
        a.fill_between(thresholds, np.array(gaussian_means) - np.array(gaussian_stds),
                       np.array(gaussian_means) + np.array(gaussian_stds), alpha=0.3, color=WONG[6])

        if i == 2:
            a.set_xlabel(r"Threshold $\nu$", fontsize=fsval+2)

        if j == 0:
            a.set_ylabel(titles[i], fontsize=fsval+2)

        if i == 0:
            a.set_title(column_titles[j], fontsize=fsval + 2)

        if i == 0 and j == 1:
            a.legend(fontsize=fsval)

plt.tight_layout()
plt.savefig("figures/minkowski.pdf", bbox_inches="tight")
plt.savefig("figures/minkowski.png", bbox_inches="tight")
plt.show()

fig, ax = plt.subplots(3, 2, figsize=(14, 12), sharex=True)
titles = ["M0 (Area)", "M1 (Perimeter)", "M2 (Genus)"]

for i, stat in enumerate(['M0', 'M1', 'M2']):
    for j, label in enumerate(['cib', 'tsz']):
        a = ax[i, j]
        train_data = results[f'{stat}_train_{label}']
        samples_data = results[f'{stat}_samples_{label}']
        gaussian_data = results[f'{stat}_gaussian_{label}']

        train_means = [x[0] for x in train_data]
        train_stds = [x[1] for x in train_data]
        samples_means = [x[0] for x in samples_data]
        samples_stds = [x[1] for x in samples_data]
        gaussian_means = [x[0] for x in gaussian_data]
        gaussian_stds = [x[1] for x in gaussian_data]

        if stat == 'M0':
            a.set_yscale('log')
        a.plot(thresholds, train_means, label='Agora', color='black', lw=2)
        a.fill_between(thresholds, np.array(train_means) - np.array(train_stds),
                       np.array(train_means) + np.array(train_stds), alpha=0.3, color='black')
        a.plot(thresholds, samples_means, label='DDPM', color='steelblue', lw=2)
        a.fill_between(thresholds, np.array(samples_means) - np.array(samples_stds),
                       np.array(samples_means) + np.array(samples_stds), alpha=0.3, color='steelblue')
        a.plot(thresholds, gaussian_means, label='Gaussian', color=WONG[6], lw=2)
        a.fill_between(thresholds, np.array(gaussian_means) - np.array(gaussian_stds),
                       np.array(gaussian_means) + np.array(gaussian_stds), alpha=0.3, color=WONG[6])

        if i == 2:
            a.set_xlabel(r"Threshold $\nu$",fontsize=fsval)
        #else:
            #a.set_xticklabels([])

        if j == 0:
            a.set_ylabel(titles[i],fontsize=fsval)
        #else:
            #a.set_yticklabels([])

        if i == 0 and j == 1:
            a.legend(fontsize=fsval)

#plt.subplots_adjust(wspace=0, hspace=0)
plt.savefig("figures/minkowski.pdf", bbox_inches="tight")
plt.savefig("figures/minkowski.png", bbox_inches="tight")
plt.show()