astro-datalab / notebooks-latest

Default set of Data Lab notebooks, by DL team and contributed by users
BSD 3-Clause "New" or "Revised" License
60 stars 49 forks source link

Introduce a function for plotting z histograms in DESI EDR notebooks #194

Open stephjuneau opened 1 year ago

stephjuneau commented 1 year ago

Suggestion from @rnikutta for a future revision of both DESI EDR tutorial notebooks. I'm marking it as an enhancement.

"This block is also very repetitive code:"

fig, axs = plt.subplots(4, 1, figsize = (9, 12))
bins = np.arange(0, 4, 0.2)

axs[0].hist(zpix_cat['z'][is_bgs], color = 'C0', bins = bins, label = f'BGS: {n_bgs} sources')
axs[0].legend(fontsize = 14)
axs[0].set_ylabel("N(z)")
axs[1].hist(zpix_cat['z'][is_lrg], color = 'C1', bins = bins, label = f'LRG: {n_lrg} sources')
axs[1].legend(fontsize = 14)
axs[1].set_ylabel("N(z)")
axs[2].hist(zpix_cat['z'][is_elg], color = 'C2', bins = bins, label = f'ELG: {n_elg} sources')
axs[2].legend(fontsize = 14)
axs[2].set_ylabel("N(z)")
axs[3].hist(zpix_cat['z'][is_qso], color = 'C3', bins = bins, label = f'QSO: {n_qso} sources')
axs[3].legend(fontsize = 14)
axs[3].set_ylabel("N(z)")
axs[3].set_xlabel("Redshift")

Why not introduce a tiny function that does the per-panel work? E.g.:

fig, axs = plt.subplots(4, 1, figsize = (9, 12))
bins = np.arange(0, 4, 0.2)

def plotaxis(axid,mask,color,cat):
    axs[axid].hist(zpix_cat['z'][mask], color = color, bins = bins, label = f'{cat}: {n_bgs} sources')
    axs[axid].legend(fontsize = 14)
    axs[axid].set_ylabel("N(z)")

plotaxis(0,is_bgs,'C0','BGS')
plotaxis(1,is_lrg,'C1','LRG')
plotaxis(2,is_elg,'C2','ELG')
plotaxis(3,is_qso,'C3','QSO')
axs[3].set_xlabel("Redshift")